Python’s mocking library is rather uncomplicated. Most of what it does is creating mock objects: veritable sponges that absorb every interaction with any code that we pass them to.

This simplicity is also surfaced in the API, especially in the main part of it — the mock.Mock constructor:

some_mock = mock.Mock(url='http://example.com')
assert some_mock.url == 'http://example.com'

Any arguments that we pass there become attributes on the resulting mock object. This is really useful when patching, because it allows us to completely specify the replacement object within a @mock.patch decorator:

@patch.object(requests, 'get', new=Mock(return_value=Mock(status_code=400)))
def test_404(self):
    self.assertRaises(NotFoundError):
        do_stuff()

You have to keep in mind, however, that the mock.Mock class also has some constructor arguments of its own. For this reason, there exists some potential for name collision: some of the Mock‘s own arguments may have the same names as the attributes we’d like to set on the mock object:

some_mock = mock.Mock(name="John Doe")  # doesn't set the `name` attribute
assert some_mock.name == "John Doe"  # blows up!

Here, the name argument is inherent to the Mock class. Its constructor will interpret it in a special way, and so it won’t set a name attribute on the resulting mock. Other possible culprits include the spec and wraps parameters, both of which have relatively common names that we may want to use as object attributes¹.

Collision avoidance

It’s trivial to fix the issue, of course:

some_mock = mock.Mock()
some_mock.name = "John Doe"
assert sock_mock.name == "John Doe"

but this approach has a downside. Creating and configuring a mock is no longer a single expression, which means we cannot use it with patchers as easily as before:

@patch.object(foo, 'bar')
def test_something(self, mock_bar):
    mock_bar.name = "John Doe"
    # (...rest of the test...)

We can either configure the mock after patching, like above, or perhaps introduce some utility functions to be called inside the @patch decorator.

The almost-there method

In any case, this is somewhat disappointing. And it is even more so when we discover that there is a method called configure_mock which looks like it was designed to solve this very issue. Its arguments are always interpreted as attributes of the mock: it has no “special” or “reserved” names. Indeed, this method is what allows us to actually write the mock setup as a single expression:

some_mock.configure_mock(name="John Doe")

Problem is, this expression returns None.

Yes, configure_mock returns nothing.
Or in other words, it doesn’t return anything.
In fact, it has no return statement whatsoever.

Most importantly, it doesn’t have the return self line that’d enable us to write this:

some_mock = mock.Mock().configure_mock(name="John Doe")

Well, that is quite a let-down.

Fixing it

But hey, this is Python! Shortcomings like that don’t necessarily mean we have to fork whole libraries. Let’s just add the missing return, shall we?

from mock import Mock as _Mock

class Mock(_Mock):
    def configure_mock(self, **kwargs):
        super(Mock, self).configure_mock(**kwargs)
        return self  # <-- there!

Whew, that was quick!

…Alright, that’s actually the whole fix, but it’s close. To complete it, we need to apply the same treatment to three more Mock classes: MagicMock, NonCallableMock, and NonCallableMagicMock.

A complete solution can be seen in this gist.

…or: The Most Hideous Hack I’ve (Almost) Done

In callee, the argument matcher library for Python that I released recently, there is this lovely TODO note for a seemingly simple feature. When using the Matching construct with a simple lambda predicate:

mock_foo.assert_called_with(Matching(lambda x: x % 2 == 0))

it would be great to see its code in the error message if the assertion fails. Right now it’s just going to say something like <Matching <function <lambda> at 0x7f5d8a06eb18>>. Provided you don’t possess a supernatural ability of dereferencing pointers in your head, this won’t give you any immediate hint as to what went wrong. Wouldn’t it be nice if it read as, say, <Matching \x: x % 2> instead?¹

So I thought: why not try and implement such a mechanism? This is Python, after all — a language where you can spawn completely new classes at runtime, walk the stack backwards (or even forward) and read the local variables, or change the behavior of the import system itself. Surely it would be possible — nay, easy — to get the source code of a short lambda function, right?

Boy, was I wrong.

Make no mistake, though: the task turned out to be absolutely doable, at least in the scope I wanted it done. But what would you think of a solution that involves not just the usual Python hackery, but also AST inspection, transformations of the source code as text, and bytecode shenanigans?…

The code, all the code, and… much more than the code

Let’s start from the beginning, though. Here’s a short lambda function, the kind of which we’d like to obtain the source code of:

is_even = lambda x: x % 2 = 0

If the documentation for Python standard library is to be believed, this should be pretty easy. In the inspect module, there is a function called no different than getsource. For our purposes, however, getsourcelines is a little more convienient, because we can easily tell when the lambda is too long:

def get_short_lambda_source(lambda_func):
    try:
        source_lines, _ = inspect.getsourcelines(lambda_func)
    except IOError:
        return None
    if len(source_lines) > 1:
        return None
    return source_lines[0].strip()

Of course if you programmed in Python for any longer period of time, you know very well that the standard docs are not to be trusted. And it’s not just that the except clause should also include TypeError, because it will be thrown when you try to pass any of the Python builtins to getsourcelines.

More important is the ambiguity of what does “source lines for an object” actually mean. “Source lines containing the object definition” would be much more accurate, and this seemingly small distinction is rather crucial here. Passing a lambda function to either getsourcelines or getsource, we’ll get its source and everything else that the returned lines included.

That’s right. Say hello to the complete is_even = assignment, and the entire assert_called_with invocation! And in case you are wondering: yes, the result will also include any end-of-line comments. No token left behind!

Trim left

Clearly this is more than we’ve bargained for. Maybe there is a way to strip away the unnecessary cruft? Python does know how to parse itself, after all: the standard ast module is a manifestation of this knowledge. Perhaps we can use it to retrieve the lambda AST node in order to turn it — and just it — back into Python code?…

def get_short_lambda_ast_node(lambda_func):
    source_text = get_short_lambda_source(lambda_func)
    if source_text:
        source_ast = ast.parse(source_text)
        return next((node for node in ast.walk(source_ast)
                     if isinstance(node, ast.Lambda)), None)

But as it turns out, getting the source text back this way is only mostly possible.

See, every substantial AST node — which is either an expression (ast.expr) or a statement (ast.stmt) — has two common attributes: lineno and col_offset. When combined, they point to a place in the original source code where the node was parsed from. This is how we can find out where to look for the definition of our lambda function.

Looks promising, right? The only problem is we don’t know when to stop looking. That’s right: nodes created by ast.parse are annotated with their start offset, but not with length nor the end offset. As a result, the best we can do when it comes to carving out the lambda source from the very first example is this:

lambda x: x % 2 == 0))

So close! Those hanging parentheses are evidently just taunting us, but how can we remove them? lambda is basically just a Python expression, so in principle it can be followed by almost anything. This is doubly true for lambdas inside the Matching construct, as they may be a part of some larger mock assertion:

mock_foo.assert_called_with(Matching(lambda x: x % 2 == 0), Integer() & GreaterThan(42))

Here, the extraneous suffix is the entirety of ), Integer() & GreaterThan(42)), quite a lot of more than just )). And that’s of course nowhere near the limit of possiblities: for one, there may be more lambdas in there, too!

Back off, slowly

It seems, however, that there is one thing those troublesome tails have in common: they aren’t syntactically valid.

Intuitively, a lambda node nested within some other syntactical constructs will have their closing fragments (e.g. )) appear somewhere after its end. Without the corresponding openings (e.g. Matching(), those fragments won’t parse.

So here’s the crazy idea. What we have is invalid Python, but only because of some unspecified number of extra characters. How about we just try and remove them, one by one, until we get something that is syntactically correct? If we are not mistaken, this will finally be our lambda and nothing else.

The fortune favors the brave, so let’s go ahead and try it:

# ... continuing get_short_lambda_source() ...

source_text = source_lines[0].strip()
lambda_node = get_short_lambda_ast_node(lambda_func)

lambda_text = source_text[lambda_node.col_offset:]
min_length = len('lambda:_')  # shortest possible lambda expression
while len(lambda_text) > min_length:
    try:
        ast.parse(lambda_text)
        return lambda_text
    except SyntaxError:
        lambda_text = lambda_text[:-1]
return None

Considering that we’re basically taking lessons from the dusty old tomes in the Restricted Section of Hogwarts library, the magic here looks quite simple. As long as there is something that can pass for a lambda definition, we try to parse it and see if it succeeds. The line that says except SyntaxError: is obviously not something for the faint of heart, but at least we are specifying what exception we anticipate catching.

And the kicker? It works. By that I mean it doesn’t return garbage results for a few obvious and not so obvious test cases, which is already more than you would normally expect from hacks of this magnitude. All the lambdas defined until this paragraph, for example, can have their source code extracted without issue.

Just one more thing

So… victory? Not quite. Astute readers may recall my promise of some bytecode arcana, and now’s the time for it.

Despite the initial success of our gradual, character dropping approach, there are cases where it doesn’t produce the correct result. Consider, for example, a lambda definition that’s nestled within a tuple²:

>>> x = lambda _: True, 0
>>> get_short_lambda_source(x[0])
lambda _: True, 0

We would of course expect the result to be lambda _: True, without a comma or zero.

Unfortunately, here’s where our earlier assumption fails rather spectacularly. The line of code extracted from AST is syntactically valid even with the extra characters. As a result, ast.parse succeeds too early and returns an incorrect definition. It should have been of a lambda contained within a tuple, but tuple is apparently what the lambda returns.

You may say that this is the sharp end of a narrow edge case, and anyone who defines functions like that deserves all the trouble they get. And sure, I wouldn’t mind if we just threw hands in the air and told them we’re simply unable to retrieve the source here. But my opinion is that it doesn’t justify serving them obviously wrong results!

A halting problem

Not if we can help it, anyway. Have a look at the expected source code and the one we’ve extracted, side by side:

lambda _: True
lambda _: True, 0

The second line isn’t just longer: it is also doing more. It isn’t just defining a lambda; it defines it, conjures up a constant 0, and then packs them both into a tuple. That’s at least two additional steps compared to the original.

Those steps have a more precise name, too: they are the bytecode instructions. Every piece of Python source is compiled to a binary bytecode before it’s executed, because the interpreter can only work with this representation. Compilation typically happens when a Python module is first imported, producing a .pyc file corresponding to its .py file. Subsequent imports will simply reuse the cached bytecode.

Moreover, any function or class object has its bytecode accessible (read-only) at runtime. There is even a dedicated data type to hold it — called simply code — with a buffer of raw bytes under one of its attributes.

Finally, the bytecode compiler itself is also available to Python programs as a built-in compile function. You don’t see it used as often as its counterparts eval and exec (which hopefully are a rare sight themselves!), but it taps into the same internal machinery of Python.

So how does it all add up? The idea is, basically, to cross-check the alleged source code of the lambda with its own bytecode. Any junk that’s still left to trim — even if syntactically valid — will surface as a divergence after compilation. Thus we can simply continue dropping characters until the bytecodes match:

lambda_text = source_text[lambda_node.col_offset:]
lambda_body_text = source_text[lambda_node.body.col_offset:]
min_length = len('lambda:_')  # shortest possible lambda expression
while len(lambda_text) > min_length:
    try:
        code = compile(lambda_body_text, '<unused filename>', 'eval')
        if len(code.co_code) == len(lambda_func.__code__.co_code):
            return lambda_text
    except SyntaxError:
        pass
    lambda_text = lambda_text[:-1]
    lambda_body_text = lambda_body_text[:-1]
return None

Okay, maybe not the exact bytes³, but stopping at the identical bytecode length is good enough a strategy. As an obvious bonus, compile will also take care of detecting syntax errors in the candidate source code, so we don’t need the ast parsing anymore.

That escalated quickly!

Believe it or not, but there aren’t any more objections to this solution, You can view it in its glorious entirety by looking at this gist.

Does it mean it is also making its cameo in the callee library?…

No, I’m afraid not.

Normally, I’m not the one to shy away from, ahem, bold solutions to tough problems. But in this case, the magnitude of hackery required is just too great, the result not satisfactory enough, the feature’s priority isn’t really all that high, and the maintenance burden it’d introduce is most likely too large.

In the end, it was great fun figuring it out: yet another example of how you can fiddle with Python to do basically anything. Still, we must not get too preoccupied with whether or not we can as to forget if we should.

Combined with dynamic typing, the lax object model of Python makes it pretty easy to write unit tests that involve mocking some parts of the tested code. Still, there’s plenty of third party libraries — usually with “mock” somewhere in the name — which promise to make it even shorter and simpler. There’s evidently been a need here that the standard library didn’t address adequately.

It changed, however, with Python 3.3. This version introduced a new unittest.mock module which essentially obsoleted all the other, third party solutions. The module, available also as a backport for Python 2.6 and above, is flexible, easy to use, and offers most if not all of the requisite functionality.

Called it, maybe?

Well, with one notable exception. If we mock a function, or any other callable object:

# (production code)
class ProductionClass(object):
    def foo(self):
        self.florb(...)

# (test code)
something = ProductionClass()
something.florb = mock.Mock()

we have very limited capabilities when it comes to verifying how they were called. Sure, we can be extremely specific:

something.florb.assert_called_with('this particular string')

or very lenient:

something.florb.assert_called_with(mock.ANY)  # works for _any_ argument

but there is virtually nothing in between. Suppose we don’t care too much what exact string the method was called with, only that it was some string that’s long enough? Things get awkward very fast then:

for posargs, kwargs in something.florb.call_args_list:
    arg = posargs[0]
    self.assertIsInstance(arg, str)
    self.assertGreaterEqual(len(arg), 16)
    break
else:
    self.fail('required call to %r not found' % (something.florb,))

Yes, this is what’s required: an actual loop through all the mock’s calls¹; unpacking tuples of positional and keyword arguments; and manual assertions that won’t even emit useful error messages when they fail.

Eww! Surely Python can do better than that?

Meet callee

Why, of course it can! Today, I’m releasing callee: a library of argument matchers compatible with unittest.mock. It enables you to write much simpler, more readable, and delightfully declarative assertions:

from callee import String, LongerOrEqual
something.florb.assert_called_with(String() & LongerOrEqual(16))

The wide array of matchers offered by callee includes numeric, string, and collection ones. And if none suits your particular needs, it’s a trivial matter to implement a custom one.

Check the library’s documentation for more examples, or jump right in to the installation guide or a full matcher reference.

In this post I’ll describe all (hopefully all!) the various ways you can specify a single dependency for a Python package.

This assumes pip is used for installation. The list of dependencies then goes either in install_requires= parameter of the setup function within setup.py, or as a separate requirements.txt file. Commonly, it will actually go in both places, with the latter being the canonical source of truth:

from setuptools import setup

with open('requirements.txt') as rf:
    setup(
        # ...
        install_requires=rf.readlines(),
    )

More details about this approach can be found in one of my previous posts.

Here, I will concentrate on the format of a single line in requirements.txt that defines a dependency. There are numerous variants that pip supports, and they are all described in excruciating detail in PEP 440. This post shall serve as a short reference on the most useful ones.

Package name (and version)

The simplest and most common option is to identify a dependency by its package name:

SQLAlchemy

This will locate it in a global index of packages, which is sometimes called a “cheese shop”. Currently, by far the most popular package registry for Python is PyPI, and pip uses it by default¹.

Without any further modifiers, pip will download and install the “current” version of the package — either the newest, or the one designated explicitly by a maintainer. This obviously makes the dependency somewhat unpredictable, for it can mean unintended upgrades that introduce breaking changes to your code.

To prevent this, you’d normally pin the dependency to an exact version²:

SQLAlchemy==0.9.10

Other comparison operators are also available:

SQLAlchemy>=0.9.10
SQLAlchemy<1.0.0

and can even be combined:

SQLAlchemy>=0.9,<1.0.0

Specs like that will make pip find the newest version that’s within given range. Assuming your dependency follows the semantic versioning scheme, this will allow you to stay on top of any minor bugfixes and improvements to an older release (0.9.x here), without the risk of accidentally upgrading to a new one (1.x) that your code is not compatible with yet.

Repository URL

Sometimes you want to live on the bleeding edge, though, and depend not just on the latest release, but the head commit to the package’s repository. This makes sense especially in large systems that are distributed among multiple repos, and where development happens in lockstep.

For those occasions, and a few others, pip can recognize direct repository URLs. They are in the format:

$VCS+$PROTOCOL://$URL@$LABEL#egg=$PACKAGE

where the $PROTOCOL part can be optional if the version control system has a default there. That’s for example the case for git, which is of course the most important VCS you’d be interested in³:

git://git.example.com/somepackage#egg=somepackage

Note that the #egg=$PACKAGE part is not a part of the $URL, and it’s only there to give a local name for the package distribution. This is what makes it possible to refer to it later via pip, if only to remove it with pip uninstall $PACKAGE. Of course, the sanest practice is to use the PyPI moniker if possible.

When no $LABEL is given, pip will use the HEAD, trunk, tip, or the equivalent default/current revision from the repo. Often though (at least in case of Git), you would also pick a branch, tag, or even a particular commit hash:

git+https://github.com/Xion/unmatcher.git@0.1.3.1#egg=unmatcher
git+ssh://github.com/You/yourpackage.git@master#egg=yourpackage
git+https://github.com/mitsuhiko/jinja2.git@5b498453b5898257b2287f14ef6c363799f1405a#egg=Jinja2

The last two options could be a good choice even with third party packages, when you don’t want to wait for a new PyPI release to get a necessary feature or an urgent bug fix.

Local filesystem

Lastly, you can ask pip to install a package from a local directory or archive. The former option is often used with the -e (--editable) flag for pip install. This installs the package in the so-called development mode, allowing you to edit its source code in-place:

$ pip install -e /home/me/Code/myotherpackage

You almost certainly don’t want to put this line in requirements.txt: you should still be pulling the other package from PyPI. But if it’s your own one — maybe a self-contained utility library used by your main program — this setup will be very helpful for making changes to it, informed by your own usage of the package.

A relatively little known feature of Python is the else block for control flow statements other than if.

If you haven’t heard about it before, you can provide such a block for both while and for loops, as well as any variant of the try statement, Its functionality is roughly analogous in both cases:

• in loops, the else block is executed if the loop didn’t exit abnormally (i.e. with break)
• in try constructs, the else block runs if no exception happened

Likely because of the unique semantics that don’t exist in other languages, neither of those constructs has been seen much in real world code. Recently, however, I’ve found they can be combined into a very pythonic pattern that’s also quite useful.

The trick

Say you have a task that won’t always succeed. Perhaps it’s a request made to a janky server, or other network operartion that’s prone to timeouts. Since failures are likely to be transient, you’d like to retry it several more times before giving up permanently.

With try/except block, you can detect those half-expected failures. With a simple loop, you can repeat the attempt for as many times as you deem feasible. Combined, they can solve the problem rather neatly:

for _ in range(MAX_RETRIES):
    try:
        # ... do stuff ...
    except SomeTransientError:
        # ... log it, sleep, etc. ...
        continue
    else:
        break
else:
    raise PermanentError()

But why?

What’s the deal with the elses here, though? Are they both necessary?

The simple answer is of course no. else after either a loop or try/except block is always a syntactic sugar. Any code that contains it can be transformed into an equivalent snippet that utilizes different techniques to achieve the same effect.

But this view isn’t very useful, for many of the essential features in any programming languages can be dismissed as superfluous using this reasoning. The real question is whether the above idiom is more readable and understandable than the alternatives.

To that, I posit, the answer is: absolutely.

Desugaring

Without the double else, this example would have to be written in a considerably more convoluted way:

retries = MAX_RETRIES
while retries > 0:
    try:
        # ... do stuff ...
        break
    except SomeTransientError:
        # ... log it, sleep, etc. ...
        retries -= 1
if retries == 0:
    raise PermanentError()

Although at first glance the difference may be minuscule, this version adds significant extra busywork the programmer has to pay careful attention to:

• The retries variable now has to be explicit, because the final conditional statement must look at its value.
• We can’t use a for loop anymore (e.g. for retries in range(MAX_RETRIES)), because we wouldn’t distinguish the “success at last try” and “retry limit exceeded” cases: they’d both result in retries equal to MAX_RETRIES - 1 after the loop¹.
• As a result, we have to remember to decrement the counter ourselves upon an error.

Additionally, the break is easy to miss amidst the actual logic within the try block, both for developer who writes the code and for any subsequent readers. An alternative is to move it outside of the try/except clause, but that in turn reintroduces continue into the except branch and further complicates the whole flow.

In short, the desugared version is more error-prone (all those off-by-ones!) and also quite inscrutable.

Python has many batteries included, but a few things are still conspicuously missing.

One of them is a standardized and convenient approach to URL manipulation, akin to the URI class in Java. There are some functions in urllib, of course (or urllib.parse in Python 3), but much like their HTTP-related comrades, they prove rather verbose and somewhat clunky.

HTTP, however, is solved by the Requests package, so you may wonder if there is some analogous package for URL operations. The answer is affirmative, and the library in question is, quite whimsically, called furl.

URL in a wrap

The sole interesting part of the furl interface is the furl class. It represents a single URL, broken down to its constituents, with properties and methods for both reading them out and replacing with new values.

Thanks to this handy (and quite obvious) abstraction, common URL operations become quite simple and self-documenting:

from furl import furl


def to_absolute(url, base):
    """If given ``url`` is a relative path,
    make it relative to the ``base``.
    """
    furled = furl(url)
    if not furled.scheme:
        return furl(base).join(url).url
    return url


def is_same_origin(*urls):
    """Check whether URLs are from the same origin (host:port)."""
    origins = set(url.netloc for url in map(furl, urls))
    return len(origins) <= 1


def get_facebook_username(profile_url):
    """Get Facebook user ID from their profile URL."""
    furled = furl(profile_url)
    if not (furled.host == 'facebook.com' or
            furled.host.endswith('.facebook.com')):
        raise ValueError("not a Facebook URL: %s" % (profile_url,))
    return furled.path.segments[-1]

# etc.

This includes the extremely prevalent, yet very harmful pattern of bulding URLs through string interpolation:

url = '%s?%s' % (BASE_URL, urlencode(query_params))

Besides looking unpythonically ugly, it’s also inflexible and error-prone. If BASE_URL gains some innate query string params ('http://example.com/?a=b'), this method will start producing completely invalid urls (with two question marks, e.g. 'http://example.com/?a=b?foo=bar').

The equivalent in furl has none of these flaws:

url = furl(BASE_URL).add(query_params).url

The full package

To see the full power of furl, I recommend having a look at its API documentation. It’s quite clear and should be very easy to use.

Celery is an asynchronous task worker that’s frequently used for background processing in Python web apps. Rather than performing a time-consuming task within the request loop, we delegate it to a queue so that a worker process can pick it up when ready. The immediate benefit is much better latency for the end user. Pros also include easier scalability, since you can adjust the number of workers to your task load.

Examples of tasks that are usually best done in the background vary from fetching data through a third-party API to sending emails, and from pushing mobile notifications to pruning the database of stale records. Anything that may take more than a couple hundred milliseconds to complete — and isn’t absolutely essential to the current HTTP request — is typically a good candidate for asynchronous processing.

In Celery, workers are just regular Python processes, often running the exact same code that’s powering the app’s web frontend¹. But unlike most of that code, they aren’t servicing any HTTP requests — they simply run some function with given arguments, both specified by whomever sent a task for execution. Indeed, those functions don’t even know who or what asked for them to be executed.

Neat, right? It is what we usually compliment as decoupling, or separation of concerns. They are valuable qualities even regardless of the UI and scaling benefits mentioned earlier.

Not quite web

But those qualities come with a trade-off. Task code is no longer a web frontend code: it doesn’t run within the comfy environment of our web framework of choice. Losing it may be quite unnerving, actually, because in a typical web application there will be many things tied directly to the HTTP request pipeline. After all, this is what web applications do — respond to HTTP requests — so it often makes perfect sense e.g. to marry the request flow with database transactions, committing or rolling those back according to HTTP status code that the app produced.

Tasks may also require a database, though, if only to assert the expected state of the world. Similar goes for memcache, a Redis instance, or basically any resource used by the frontend code. Alas, it’s quite possible the very reason we delegate work to a task is to shift lengthy interactions with those external systems away from the UI. Obviously, we’re going to need them for that!

Fake a request

So one way or another, our tasks will most likely need some initialization and/or cleanup code. And since it’s probably the same code that most HTTP request handlers require and use already, why not just pretend we’re handling a request after all?

In Flask, we can pull that off rather easily. The test_request_context method is conveniently provided to allow for faking the request context — that is, an execution environment for HTTP handlers. Like the name suggest, it is used mostly for testing, but there is nothing stopping us from using it in tasks run by Celery.

We probably don’t want to call it directly, though. What would be better is to have Celery prepare the context first, and then run the task code as if it was an HTTP handler. For even better results, the context would preserve information extracted from the actual HTTP request, one that has sent the task for execution. Moving some work to the background would then be a trivial matter, for both the task and the original handler would operate within the same environment.

Convenient? I believe so. And as I’ll demonstrate next, it isn’t very complicated to implement either.

Wrap the decorator?

At least one piece of the solution should stand out as pretty obvious. Since our intention is to wrap the task’s code in some additional packaging — the request context — it seems fairly natural to write our own @task decorator:

import functools

from myapp import app, celery


def task(**kwargs):
    """Decorator function to apply to Celery tasks.
    Executes the actual task inside Flask's test request context.
    """
    def decorator(func):
        """Actual decorator."""
        @celery.task(**kwargs)
        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            with app.test_request_context():
                return func(*args, **kwargs)

        return wrapped

    return decorator

Here, app is the Flask application, and celery is the Celery object that’s often configured alongside it.

The Task class

While this technique will give us a request context inside the task code, it won’t be the request context from which the task has been sent for execution. To replicate that context correctly, we need additional support on the sender side.

Internally, Celery is converting every function we annotate with @task decorator into a subclass of Celery.app.task.Task. This process can be customized, for example by providing an explicit base= parameter to @task which specifies a custom Task subclass to use in place of the default one. Within that subclass, we’re free to override any functionality that we need to.

In our case, we’ll scrape the current request context from Flask for all relevant information, and later use it to recreate the context in the task code.

But before we get to that, let’s just create the Task subclass and move the above execution logic to it. This way, users won’t have to use a completely new @task decorator which would needlessly couple them to a specific Celery app instance:

from celery import Task

class RequestContextTask(Task):
    """Base class for tasks that run inside a Flask request context."""
    abstract = True

    def __call__(self, *args, **kwargs):
        with app.test_request_context():
            return super(RequestContextTask, self).__call__(*args, **kwargs)

Instead, they can either set this class as base= for a specific task:

@celery.task(base=RequestContextTask)
def sync_user_data(user_id):
    # ...

or make it into new default for all their tasks:

celery = Celery(...)
celery.Task = RequestContextTask

Invocation patterns

When the frontend asks for a task to be executed, it most often uses the Task.delay method. It will package a payload contaning task arguments, and send it off through a broker — usually an AMQP-based queue, such as RabbitMQ — so that a Celery worker can pick it up and actually execute.

But there are other means of task invocation. We can even run it “in place”, locally and synchronously, which is especially useful for various testing scenarios. Lastly, a task can also be retried from within its own code, terminating its current run and scheduling another attempt for some future date.

Obviously, for the RequestContextTask to be useful, it needs to behave correctly in every situation. Therefore we need to cover all the entry points I’ve mentioned — the asynchronous call, a synchronous invocation, and a task retry:

class RequestContextTask(Task):
    # ...

    def apply_async(self, args=None, kwargs=None, **rest):
        self._include_context(kwargs)
        return super(RequestContextTask, self) \
            .apply_async(args, kwargs, **rest)

    def apply(self, args=None, kwargs=None, **rest):
        self._include_context(kwargs)
        return super(RequestContextTask, self) \
            .apply(args, kwargs, **rest)

    def retry(self, args=None, kwargs=None, **rest):
        self._include_context(kwargs)
        return super(RequestContextTask, self) \
            .retry(args, kwargs, **rest)

Note that Task.apply_async is being called internally by Task.delay, so it’s only that first method that we have to override.

Context in a box

As you can deduce right away, the Flask-related magic is meant to go in the _include_context method. The idea is to prepare arguments for the eventual invocation of Flask.test_request_context, and pass them through an extra task parameter. Those arguments are relatively uncomplicated: they are just a medley of various pieces of information that we can easily obtain from the Flask’s request object:

from flask import has_request_context, request


class RequestContextTask(Task):
    CONTEXT_ARG_NAME = '_flask_request_context'

    # ...

    def _include_request_context(self, kwargs):
        """Includes all the information about current HTTP request context
        as an additional argument to the task.
        """
        if not has_request_context():
            return

        context = {
            'path': request.path,
            'base_url': request.url_root,
            'method': request.method,
            'headers': dict(request.headers),
        }
        if '?' in request.url:
            context['query_string'] = request.url[(request.url.find('?') + 1):]

        kwargs[self.CONTEXT_ARG_NAME] = context

On the worker side, we simply unpack them and recreate the context:

from flask import make_response

class RequestContextTask(Task):
    # ...
    def __call__(self, *args, **kwargs):
        call = lambda: super(RequestContextTask, self).__call__(*args, **kwargs)

        context = kwargs.pop(self.CONTEXT_ARG_NAME, None)
        if context is None or has_request_context():
            return call()

        with app.test_request_context(**context):
            result = call()
            app.process_response(make_response(result or ''))

        return result

The only tricky part is calling Flask.process_response at the end. We need to do that for the @after_request hooks to execute correctly. This is quite crucial, because those hooks are where you’d normally put important cleanup code, like commit/rollback of the database transaction.

Complete solution

To see how all those code snippets fit together, see this gist. You may also want to have a look at the article on Celery integration in Flask docs for some tips on how to integrate it with your own project.

Little over two years ago, I was writing about my skepticism towards the addition of enum types to Python. My stance has changed somewhat since then, and I now think there is some non-trivial value that enums can add to Python code. It becomes especially apparent in the circumstances involving any kind of persistence, from pickling to databases.

I will elaborate on that through the rest of this post.

Revision

First, let’s recap (or perhaps introduce) the important facts about enums in Python. An enum, or an enumeration type, is a special class that inherits from enum.Enum¹ and defines one or more constants:

from enum import Enum

class Cutlery(Enum):
    knife = 'knife'
    fork = 'fork'
    spoon = 'spoon'

Their values are then converted into singleton objects with distinct identity. They available as attributes of the enumeration type:

>>> Cutlery.fork
<Cutlery.fork: 'fork'>

As a shorthand for when the values themselves are not important, the Enum class can be directly invoked with constant names as strings:

Cutlery = Enum('Cutlery', ['knife', 'fork', 'spoon'])

Resulting class offers some useful API that we’d expect from an enumeration type in any programming language:

>>> list(Cutlery)
[<Cutlery.knife: 'knife'>, <Cutlery.fork: 'fork'>, <Cutlery.spoon: 'spoon'>]
>>> [c.value for c in Cutlery]
['knife', 'fork', 'spoon']
>> Cutlery.knife in Cutlery
True
>>> Cutlery('spoon')
<Cutlery.spoon: 'spoon'>
>>> Cutlery('spork')
Traceback (most recent call last):
# (...)
ValueError: 'spork' is not a valid Cutlery

How it was done previously

Of course, there is nothing really groundbreaking about assigning some values to “constants”. It’s been done like that since times immemorial, and quite often those constants have been grouped within finite sets.

Here’s an example of a comment model in some imaginary ORM, complete with a status “enumeration”:

class Comment(Model):
    APPROVED = 'approved'
    REJECTED = 'rejected'
    IN_REVIEW = 'in_review'
    STATES = frozenset([APPROVED, REJECTED, IN_REVIEW])

    author = StringProperty()
    text = TextProperty()
    state = String(choices=STATES)

# usage
comment = Comment(author="Xion", text="Boo")
comment.state = Comment.IN_REVIEW
comment.save()

Converting it to use an Enum is relatively straightforward:

class Comment(Model):
    class State(Enum):
        approved = 'approved'
        rejected = 'rejected'
        in_review = 'in_review'

    author = StringProperty()
    text = TextProperty()
    state = StringProperty(choices=[s.value for s in State])

comment = Comment(author="Xion", text="Meh.")
comment.state = Comment.State.approved.value
comment.save()

It is not apparent, though, if there is any immediate benefit of doing so. True, we no longer need to define the STATES set explicitly, but saving on this bit of boilerplate is balanced out by having to access the enum’s value when assigning it to a string property.

All in all, it seems like a wash — though at least we’ve established that enums are no worse than their alternatives :)

Enums are interoperable

Obviously, this doesn’t exactly sound like high praise. Thing is, we haven’t really replaced the previous solution completely. Remnants of it still linger in the way the state property is declared. Even though it’s supposed to hold enum constants, it is defined as string, which at this point is more of an implementation detail of how those constants are serialized.

What were really need here is a kind of EnumProperty that’d allow us to work solely with enum objects. Before the introduction of a standard Enum base, however, there was little incentive for ORMs and other similiar libraries to provide such a functionality. But now, it makes much more sense to support enums as first-class citizens, at least for data exchange and serialization, because users can be expected to already prefer the standard Enum in their own code.

Thus, the previous example changes into something along these lines:

class Comment(Model):
    # ...
    state = EnumProperty(State)

comment = Comment(author="Xion", text="Yay!")
comment.state = Comment.State.approved  # no .value!
comment.save()

Details of EnumProperty, or some equivalent construct, are of course specific to any given data management library. In SQLAlchemy, for example, a custom column type can handle the necessary serialization and deserialization between Python code and SQL queries, allowing you to define your models like this²:

class Comment(Model):
    class State(Enum):
        # ...

    author = Column(String(255))
    text = Column(Text)
    state = Column(EnumType(State, impl=String(32)))

# usage like above

In any case, the point is to have Python code operate solely on enum objects, while the persistence layer takes care of converting between them and their serializable values.

Enums are extensible

The other advantage Enums have over loose collections of constants is that they are proper types. Like all user-defined types in Python, they can have additional methods and properties defined on their instances, or even on the type itself.

Although this capability is (by default) not as powerful as e.g. in Java — where each enum constant can override a method in its own way — it can nevertheless be quite convenient at times. Typical use cases include constant classification:

class Direction(Enum):
    left = 37
    up = 38
    right = 39
    down = 40

    @property
    def is_horizontal(self):
        return self in (Direction.left, Direction.right)

    @property
    def is_vertical(self):
        return self in (Direction.down, Direction.up)

and conversion:

    def as_vector(self):
        return {
            Direction.left: (-1, 0),
            Direction.up: (0, -1),
            Direction.right: (1, 0),
            Direction.down: (0, 1),
        }.get(self)

For the latter, it would be handy to have the Java’s ability to attach additional data to an enum constant. As it turns out, Python supports this feature natively in a very similar way. We simply have to override enum’s __new__ method to parse out any extra values from the initializer and turn them into attributes of the enum instance:

class Direction(Enum):
    left = 37, (-1, 0)
    up = 38, (0, -1)
    right = 39, (1, 0)
    down = 40, (0, 1)

    def __new__(cls, keycode, vector):
        obj = object.__new__(cls)
        obj._value_ = keycode
        obj.vector = vector
        return obj

It’s possible, in fact, to insert any arbitrary computation here that yields the final _value_ of an enum constant³. This trick can be used to, for example, construct enums that automatically number themselves.

Finally, we can add static methods, class methods, or class properties to the Enum subclass, just like we would do with any other class:

class MyEnum(Enum):
    # ...

    @classproperty
    def __values__(cls):
        return [m.value for m in cls]

Enums just are

All these things are possible primarly because of the most notable aspect of Python enums: their existence as an explicit concept. A syntactically unorganized bunch of constants cannot offer half of the highlighted features because there is nowhere to pin them on.

For that reason alone, using enums as an explicit — rather than implicit — pattern seems worthwhile. The one benefit we’re certain to reap is better code structure through separation of important concepts.