Arguments to Python generator functions

Posted on Tue 14 March 2017 in Code • Tagged with Python, generators, functions, arguments, closuresLeave a comment

In Python, a generator function is one that contains a yield statement inside the function body. Although this language construct has many fascinating use cases (PDF), the most common one is creating concise and readable iterators.

A typical case

Consider, for example, this simple function:

def multiples(of):
    """Yields all multiples of given integer."""
    x = of
    while True:
        yield x
        x += of

which creates an (infinite) iterator over all multiples of given integer. A sample of its output looks like this:

>>> from itertools import islice
>>> list(islice(multiples(of=5), 10))
[5, 10, 15, 20, 25, 30, 35, 40, 45, 50]

If you were to replicate in a language such as Java or Rust — neither of which supports an equivalent of yield — you’d end up writing an iterator class. Python also has them, of course:

class Multiples(object):
    """Yields all multiples of given integer."""

    def __init__(self, of):
        self.of = of
        self.current = 0

    def __iter__(self):
        return self

    def next(self):
        self.current += self.of
        return self.current

    __next__ = next  # Python 3

but they are usually not the first choice1.

It’s also pretty easy to see why: they require explicit bookkeeping of any auxiliary state between iterations. Perhaps it’s not too much to ask for a trivial walk over integers, but it can get quite tricky if we were to iterate over recursive data structures, like trees or graphs. In yield-based generators, this isn’t a problem, because the state is stored within local variables on the coroutine stack.

Lazy!

It’s important to remember, however, that generator functions behave differently than regular functions do, even if the surface appearance often says otherwise.

The difference I wanted to explore in this post becomes apparent when we add some argument checking to the initial example:

def multiples(of):
    """Yields all multiples of given integer."""
    if of < 0:
        raise ValueError("expected a natural number, got %r" % (of,))

    x = of
    while True:
        yield x
        x += of

With that if in place, passing a negative number shall result in an exception. Yet when we attempt to do just that, it will seem as if nothing is happening:

>>> m = multiples(-10)
>>>

And to a certain degree, this is pretty much correct. Simply calling a generator function does comparatively little, and doesn’t actually execute any of its code! Instead, we get back a generator object:

>>> m
<generator object multiples at 0x10f0ceb40>

which is essentially a built-in analogue to the Multiples iterator instance. Commonly, it is said that both generator functions and iterator classes are lazy: they only do work when we asked (i.e. iterated over).

Getting eager

Oftentimes, this is perfectly okay. The laziness of generators is in fact one of their great strengths, which is particularly evident in the immense usefulness of theitertools module.

On the other hand, however, delaying argument checks and similar operations until later may hamper debugging. The classic engineering principle of failing fast applies here very fittingly: any errors should be signaled immediately. In Python, this means raising exceptions as soon as problems are detected.

Fortunately, it is possible to reconcile the benefits of laziness with (more) defensive programming. We can make the generator functions only a little more eager, just enough to verify the correctness of their arguments.

The trick is simple. We shall extract an inner generator function and only call it after we have checked the arguments:

def multiples(of):
    """Yields all multiples of given integer."""
    if of < 0:
        raise ValueError("expected a natural number, got %r" % (of,))

    def multiples():
        x = of
        while True:
            yield x
            x += of

    return multiples()

From the caller’s point of view, nothing has changed in the typical case:

>>> multiples(10)
<generator object multiples at 0x110579190>

but if we try to make an incorrect invocation now, the problem is detected immediately:

>>> multiples(-5)
Traceback (most recent call last):
  File "<pyshell#2>", line 1, in <module>
    multiples(of=-5)
  File "<pyshell#0>", line 4, in multiples
    raise ValueError("expected a natural number, got %r" % (of,))
ValueError: expected a natural number, got -5

Pretty neat, especially for something that requires only two lines of code!

The last (micro)optimization

Indeed, we didn’t even have to pass the arguments to the inner (generator) function, because they are already captured by the closure.

Unfortunately, this also has a slight performance cost. A captured variable (also known as a cell variable) is stored on the function object itself, so Python has to emit a different bytecode instruction (LOAD_DEREF) that involves an extra pointer dereference. Normally, this is not a problem, but in a tight generator loop it can make a difference.

We can eliminate this extra work2 by passing the parameters explicitly:

    # (snip)

    def multiples(of):
        x = of
        while True:
            yield x
            x += of

    return multiples(of)

This turns them into local variables of the inner function, replacing the LOAD_DEREF instructions with (aptly named) LOAD_FAST ones.


  1. Technically, the Multiples class is here is both an iterator (because it has the next/__next__ methods) and iterable (because it has __iter__ method that returns an iterator, which happens to be the same object). This is common feature of iterators that are not associated with any collection, like the ones defined in the built-in itertools module

  2. Note that if you engage in this kind of microoptimizations, I’d assume you have already changed your global lookup into local ones :) 

Continue reading

Optional arguments in Rust 1.12

Posted on Thu 29 September 2016 in Code • Tagged with Rust, arguments, parameters, functionsLeave a comment

Today’s announcement of Rust 1.12 contains, among other things, this innocous little tidbit:

Option implements From for its contained type

If you’re not very familiar with it, From is a basic converstion trait which any Rust type can implement. By doing so, it defines how to create its values from some other type — hence its name.

Perhaps the most widespread application of this trait (and its from method) is allocating owned String objects from literal str values:

let hello = String::from("Hello, world!");

What the change above means is that we can do similar thing with the Option type:

let maybe_int = Option::from(42);

At a first glance, this doesn’t look like a big deal at all. For one, this syntax is much more wordy than the traditional Some(42), so it’s not very clear what benefits it offers.

But this first impression is rather deceptive. In many cases, this change can actually reduce the number of times we have to type Some(x), allowing us to replace it with just x. That’s because this new impl brings Rust quite a bit closer to having optional function arguments as a first class feature in the language.

Until now, a function defined like this:

fn maybe_plus_5(x: Option<i32>) -> i32 {
    x.unwrap_or(0) + 5
}

was the closest Rust had to default argument values. While this works perfectly — and is bolstered by compile-time checks! — callers are unfortunately required to build the Option objects manually:

let _ = maybe_plus_5(Some(42));  // OK
let _ = maybe_plus_5(None);      // OK
let _ = maybe_plus_5(42);        // error!

After Option<T> implements From<T>, however, this can change for the better. Much better, in fact, for the last line above can be made valid. All that is necessary is to take advantage of this new impl in the function definition:

fn maybe_plus_5<T>(x: T) -> i32 where Option<i32>: From<T> {
    Option::from(x).unwrap_or(0) + 5
}

Unfortunately, this results in quite a bit of complexity, up to and including the where clause: a telltale sign of convoluted, generic code. Still, this trade-off may be well worth it, as a function defined once can be called many times throughout the code base, and possibly across multiple crates if it’s a part of the public API.

But we can do better than this. Indeed, using the From trait to constrain argument types is just complicating things for no good reason. What we should so instead is use the symmetrical trait, Into, and take advantage of its standard impl:

impl<T, U> Into<U> for T where U: From<T>

Once we translate it to the Option case (now that Option<T> implements From<T>), we can switch the trait bounds around and get rid of the where clause completely:

fn maybe_plus_5<T: Into<Option<i32>>>(x: T) -> i32 {
    x.into().unwrap_or(0) + 5
}

As a small bonus, the function body has also gotten a little simpler.


So, should you go wild and change all your functions taking Optionals to look like this?… Well, technically you can, although the benefits may not outweigh the downsides for small, private functions that are called infrequently.

On the other hand, if you can afford to only support Rust 1.12 and up, this technique can make it much more pleasant to use the external API of your crates.

What’s best is the full backward compatibility with any callers that still pass Some(x): for them, the old syntax will continue to work exactly like before. Also note that the Rust compiler is smart about eliding the no-op conversion calls like the Into::into above, so you shouldn’t observe any changes in the performance department either.

And who knows, maybe at some point Rust makes the final leap, and allows skipping the Nones?…

Continue reading

Source code of a Python lambda

Posted on Tue 19 April 2016 in Code • Tagged with Python, functions, AST, bytecodeLeave a comment

…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?1

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 tuple2:

>>> 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 bytes3, 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.


  1. Backslash (\) is how lambda functions are denoted in Haskell. We want to be short and sweet, so it feels like a natural choice. 

  2. This isn’t an actual snippet from a Python REPL, because inspect.getsourcelines requires the object to be defined in a .py file. 

  3. Why we won’t always get an identical bytecode? The short answer is that some instructions may be swapped for their approximate equivalents.
    The long answer is that with compile, we aren’t able to replicate the exact closure environment of the original lambda. When a function refers to an free variable (like foo in lambda x: x + foo), it is its closure where the value for that variable comes from. For ad-hoc lambdas, this is typically the local scope of its outer function.
    Code produced by compile, however, isn’t associated with any such local scope. All free names are thus assumed to refer to global variables. Because Python uses different bytecode instructions for referencing local and global names (LOAD_FAST vs LOAD_GLOBAL), the result of compile may differ from a piece of bytecode produced in the regular manner. 

Continue reading

Decorated functions in Go

Posted on Fri 04 September 2015 in Code • Tagged with Go, functions, decoratorsLeave a comment

One of the downsides of working with a simple and minimalistic programming language such as Go is that abstracting for code reuse is often rather challenging. What’s brought up most often in this context is Go’s lack of support for generics (higher order types).

But here I wanted to talk about a different mechanism that — unlike generics — can be transplanted to Go quite successfully: function decorators.

Canonical example

Decorators are most commonly found in Python (where they can be applied not just to functions, but also classes), but have syntactical analogues in Java (annotations), C# (attributes), and a few other languages.

A decorator looks like a modifier adorning a function definition, distinguished by its initial @ sign:

@app.route('/home')
@login_required
def home():
    return render_template('user/home.html')

Even if you’re not familiar with the particular web framework this request handler is meant for — or indeed, the Python language itself — it shouldn’t be too difficult to figure out what purpose the two decorators serve.

What’s really important is that their goals are neatly separated from essential logic of the request handler:

  • it doesn’t have to explicitly check if the user is logged in — it has the @login_required decorator
  • it doesn’t have to be separately registered in some centalized URL routing choke point — it has the @app.route decorator instead

Such separation of concerns is a desirable quality in software, because it makes it easier to reason about different aspects of the system.

Under the hood

Behind this synctactic sugar and lofty phrasing, the code above is still perfectly equivalent to its undecorated version:

def home():
    if not current_user.is_authenticated:
        abort(403, "Login Required")
    return render_template('user/home.html')

app.add_url_route('/home', view_func=home)

No source code transformations take place, either: @login_required doesn’t actually “inject” the if statement at the beginning of home function. What happens instead is that it takes the whole function and wraps it in a new one which also contains the crucial check:

def login_required(func):
    """Grossly simplified version of a decorator enforcing user login."""
    def wrapped(*args, **kwargs):
        if not current_user.is_authenticated:
            abort(403, "Login Required")
        return func(*args, **kwargs)

    return wrapped

Given this definition, decorating home simply means passing it to the login_required function, and calling its result a new home:

def home():
    return render_template('user/home.html')

home = login_required(home)

As you’ve probably guessed, the @ syntax is nothing else than a syntatic sugar for the above. But how sweet a sugar it is!

(The @app.route decorator is simultaneously simpler and more complicated. I recommend having a direct look at its source code for more insight).

What gophers can learn

Unlike Python, in Go, we are out of luck when it comes to dedicated language support for decorators. However, the general principle still applies: we can decorate functions by writing — ahem — decorator functions.

Consider, for example, a trivial net/http request handler:

func hello(w http.ResponseWriter, r *http.Request) {
    fmt.Fprint(w, "Hello, world!\n")
}

Any real web application will have many similar handlers. They are usually wired to their corresponding URL paths during program startup:

func main() {
    http.HandleFunc("/hello", hello)
    // ...
}

This approach to routing configuration provides an opportunity to decorate some (or all) of the request handlers with any auxiliary functionality that the application may require.

Handler in, handler out

To do that, though, we need to write the necessary decorator functions first. For a simplest possible example, let’s create one that automatically fills in the Server: header of HTTP response with the name and version of our web application:

func WithServerHeader(h http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        w.Header().Set("Server", "HelloServer v0.0.1")
        h(w, r)
    }
}

Using it is then a simple matter of wrapping the original handler in a decorator call:

http.HandleFunc("/hello", WithServerHeader(hello))

As it both accepts and returns a http.HandlerFunc — a function type matching request handlers’ signatures — we don’t need anything else to be able to wrap our handlers in as many decorators as necessary:

http.HandleFunc("/", WithServerHeader(index))
http.HandleFunc("/home", WithServerHeader(LoginRequired(home)))
http.HandleFunc("/api/share", WithServerHeader(LoginRequired(CsrfProtected(api.Share))))

It’s possible simply by the means of ordinary function composition.

Going meta

Looking at the last line of the example above — which is made-up but could very well come from actual production code — you probably can’t help but notice that stacking function calls like that creates a somewhat messy amalgamation of parentheses. Beyond three or four items, a decorator chain such as this one becomes rather unwieldy to follow.

If we’re concerned about readability, it’s possible to alleviate the issue by taking another step up the abstraction ladder. Rather than composing the decorators explicitly, we can write a function that does it for us:

type HttpHandlerDecorator func(http.HandlerFunc) http.HandlerFunc

func Handler(h http.HandlerFunc, decors ...HttpHandlerDecorator) http.HandlerFunc {
    for i := range decors {
        d := decors[len(decors) - 1 - i]  // iterate in reverse
        h = d(h)
    }
    return h
}

and thus eliminate all the parentheses:

http.HandleFunc("/api/share", Handler(api.Share,
    WithServerHeader, LoginRequired, CsrfProtected))

This technique proves especially valuable if the decorators themselves are parameterized:

http.HandleFunc("/api/notifications", Handler(api.Notifications,
    WithServerHeader, LoginRequired, Cached(time.Duration(5)*time.Minute)))

In practice, you may find it valuable to wrap the complete http.HandleFunc call to accept a request handler along with its decorator chain, or the equivalent of such a call in any of the numerous Go web frameworks.

Continue reading