Dan Gittik Let’s talk

Writing/Metaprogramming: An Adventure

Chapter 8 · 24 min read

A Metaprogramming Layer

Honestly, though, in most cases we're not really going to augment the Python interpreter – definitely not for something as minute as adding a regex comparison operator. In fact, I'd argue we didn't really add one anyway, precisely because it's such a huge undertaking – instead, we resorted to using ast to change one code into another. But that's a pretty cool idea in and of itself, isn't it? One, I think, that warrants an exploration of its own.

Auto-Logging: with Tracers

Let's take logging as an example: adding prints to code can be useful, and we can always just do it explicitly. But as we've seen, mixing such a cross-cutting concern into the main flow tends to complicate it: 

>>> def add(x, y):
...     print(f"{x=}")
...     print(f"{y=}")
...     result = x + y
...     print(f"{result=}")
...     return result

That's 6 lines instead of 2 (or even 1) – for one of the simplest scenarios imaginable!

Of course, logging a function's invocation and return value can be achieved in many other ways, so let's go over the techniques we covered and see which ones can be of help. One thing we could do is use a tracer: 

>>> def log(frame, event, arg):
...     if event == "call":
...         print(f"enter {format_call(frame)}")
...     if event == "return":
...         print(f"leave {format_call(frame)}: {arg!r}")
...     return log

>>> def format_call(frame):
...     name = frame.f_code.co_name
...     params = frame.f_code.co_varnames[:frame.f_code.co_argcount]
...     args = [frame.f_locals[param] for param in params]
...     return f"{name}({', '.join(repr(arg) for arg in args)})"

>>> import sys
>>> sys.settrace(log)

This gives us the same logs without us having to add a single line: 

>>> def add(x, y):
...     return x + y

>>> add(1, 2)
enter add(1, 2)
leave add(1, 2): 3
3

But it's probably an overkill: tracers are rather slow, and it's not so easy to work with frame objects, either – supporting for keyword arguments, for example, or detecting when a function exits with an error, aren't trivial fixes to make.

Auto-Logging: with Decorators

Recall that when we discussed aspect-oriented programming, we saw that all it takes to address such common problems as logging is plugging our code into a framework, which can be implemented in a much more straightforward fashion – like, by way of a decorator: 

>>> def log(function):
...     def wrapper(*args, **kwargs):
...         call = format_call(function, *args, **kwargs)
...         try:
...             print(f"enter {call}")
...             result = function(*args, **kwargs)
...             print(f"leave {call}: {result!r}")
...             return result
...         except Exception as error:
...             print(f"leave {call} on error: {error}")
...             raise
...     return wrapper

>>> def format_call(function, *args, **kwargs):
...     args = [repr(arg) for arg in args]
...     args.extend(f"{key}={value!r}" for key, value in kwargs.items())
...     return f"{function.__name__}({', '.join(args)})"

This achieves the same result with negligible overhead: 

>>> @log
... def add(x, y):
...     return x + y

>>> add(1, 2)
enter add(1, 2)
leave add(1, 2): 3
3

And incidentally, supports keyword arguments and errors just as well. Its main drawback is the registration method: the fact that we have to add @log on top of every single invocation we want to log means there's no quick and easy way to "just apply" it to on an entire project – not unless we dive into its source code and edit each function accordingly.

Auto-Logging: with Dynamic Execution

Or is there? We've also seen that one code can easily execute another, so we can write a program like with_logs.py that will load a module and auto-decorate all the callables in its namespace, effectively changing the pointcut from "anything with @log on top" to "anything at all". There's just one problem – suppose we put this code in example.py: 

def add(x, y):
    return x + y
    
add(1, 2)

If python with_logs.py example.py were to load the code and decorate it, add(1, 2) would be invoked at load time, before we'd have a chance to apply @log to the add function.

So, let's do a trick: instead of executing the code into a regular dictionary, let's do it into one that decorates callables as soon as they're defined, which we can intercept in __setitem__: 

>>> class LoggedNamespace(dict):
...     def __setitem__(self, name, value):
...         if callable(value):
...             value = log(value)
...         super().__setitem__(name, value)

>>> def run_with_logs(path):
...     code = open(path).read()
...     namespace = LoggedNamespace()
...     exec(code, namespace)

>>> run_with_logs("example.py")
enter add(1, 2)
leave add(1, 2): 3

Still, there are some drawbacks to this approach: for one, it won't catch nested or dynamically defined functions, since LoggedNamespace only replaces the global scope. Moreover, we've introduced what's effectively a level of virtualization: it's no longer our program being executed – it's with_logs.py executing our program. If anything goes wrong, this setup is considerably harder to debug: the executed code doesn't include any logging, while the executor doesn't include any of the original logic; it's only in their dynamic interplay that the two form a coherent flow. In other words, we don't see the full picture in either source file, or even in both – it only emerges at runtime. It doesn't matter much in our case, since it's exceedingly simple, but in a real application mixing logics this way has a cost: not only does it increase the complexity, but it does so in a liminal space that's much more difficult to observe and reason about.

Auto-Logging: with AST

In some sense, augmenting functionality dynamically is akin to an interpreter, which operates on our behalf; what we'd like instead, in this analogy, is more of a compiler: something that inspects our code, augments it, and spits out a result – a flat output, so to speak, that we can look and poke at without any invisible, virtual dimensions involved. Maybe it's not a compiler per say, since it doesn't lower our code into a more basic language; and not exactly a transpiler either, since it doesn't translate between languages; but however we call it, its should rewrite the original code to include log statements.

So how do we do it? Let's start with return statements, since it's easier: all we want is to swap them with a version that prints the return value before returning – so, rewrite return <expression> into something like: 

_ = <expression>
print(_)
return _

Let's start by comparing their ASTs: 

>>> import ast
>>> ast.dump(ast.parse("return 1 + 2"), indent=2)
Module(
  body=[
    Return(
      value=BinOp(
        left=Constant(value=1),
        op=Add(),
        right=Constant(value=2)))],
  type_ignores=[])
>>> ast.dump(ast.parse("_ = 1 + 2; print(_); return _"), indent=2)
Module(
  body=[
    Assign(
      targets=[
        Name(id='_', ctx=Store())],
      value=BinOp(
        left=Constant(value=1),
        op=Add(),
        right=Constant(value=2))),
    Expr(
      value=Call(
        func=Name(id='print', ctx=Load()),
        args=[
          Name(id='_', ctx=Load())],
        keywords=[])),
    Return(
      value=Name(id='_', ctx=Load()))],
  type_ignores=[])

What we need, then, is a transformer that visits Return nodes, and replaces each with three statements: an assignment of the returned expression, followed by a call to print it, followed by a return statement of the assigned name. We've done things like this before – the only catch is that this time, we need to call self.generic_visit(node), which recursively transforms "the rest of" our node, so that we can use the fully transformed node.value in our assignment: 

>>> class ReturnLogger(ast.NodeTransformer):
...     def visit_Return(self, node):
...         self.generic_visit(node)
...         assign = ast.Assign(
...             targets=[ast.Name(id="_", ctx=ast.Store())],
...             value=node.value,
...         )
...         print_ = ast.Expr(
...             value=ast.Call(
...                 func=ast.Name(id="print", ctx=ast.Load()),
...                 args=[ast.Name(id="_", ctx=ast.Load())],
...                 keywords=[],
...             ),
...         )
...         return_ = ast.Return(
...             value=ast.Name(id="_", ctx=ast.Load()),
...         )
...         return [assign, print_, return_]

>>> def add_logs(source, target):
...     code = open(source).read()
...     tree = ast.parse(code)
...     tree = ReturnLogger().visit(tree)
...     ast.fix_missing_locations(tree)
...     open(target, "w").write(ast.unparse(tree))

>>> add_logs("example.py", "logged_example.py")

I'm also using fix_missing_locations to fill line numbers and column offsets automatically; this might introduce some discrepencies, but it's good enough for our purposes. As for the generated logged_example.py: 

def add(x, y):
    _ = x + y
    print(_)
    return _
add(1, 2)

It's exactly what we wanted – and running it prints 3, which it didn't before.

A Metaprogramming Layer

Still, if we want to recreate the nice, informative messages we had before, we need to figure out how to get the function name and arguments, as well as a way to inject the enter message at the function's beginning. Since transformers operate recursively, we know that visit_Return happens as part of a FunctionDef – that is, if we implement the latter, we can set the name and parameters list there and call self.generic_visit(node). As we recurse "into" each node's return-statements, we can count on those values being available for reconstruction: 

>>> class ReturnLogger(ast.NodeTransformer):
...     def visit_FunctionDef(self, node):
...         self.function_name = node.name
...         self.function_params = [arg.arg for arg in node.args.args]
...         self.generic_visit(node)
...         return node
...     def visit_Return(self, node):
...         # … same as before, except:
...         print_ = ast.Expr(
...             value=ast.Call(
...                 func=ast.Name(id="print", ctx=ast.Load()),
...                 args=[self.format_call("leave")],
...                 keywords=[],
...             ),
...         )

Only... how do we format the call? We've collected the function name and parameters at definition time, great; but what we want to print is their arguments, at invocation time, which is a whole other matter. Well, how would we do it if we were writing the code manually? 

>>> def add(x, y):
...     _ = x + y
...     print(f"leave add({x!r}, {y!r}): {_!r}")
...     return _

So in fact, we don't need the resolved values – since we're replacing code before it's being executed, we only need to regenerate it in such a way that once it is, those values will be resolved; in other words, we need an f-string. The way these are implemented in AST is a bit tedious, because each part has to be addressed separately – but it is what it is: 

>>> class ReturnLogger(ast.NodeTransformer):
...     # … same as before
...     def format_call(self, action):
...         parts = [
...             ast.Constant(value=f"{action} {self.function_name}("),
...         ]
...         for n, param in enumerate(self.function_params, 1):
...             parts.append(ast.FormattedValue(
...                 value=ast.Name(id=param, ctx=ast.Load()),
...                 conversion=114, # that's "!r"
...             ))
...             is_last = n == len(self.function_params)
...             if not is_last:
...                 parts.append(ast.Constant(value=", "))
...         parts.append(ast.Constant(value=")"))
...         if action == "leave":
...             parts.append(ast.Constant(value=": "))
...             parts.append(ast.FormattedValue(
...                 value=ast.Name(id="_", ctx=ast.Load()),
...                 conversion=114,
...             ))
...         return ast.JoinedStr(parts) # that's an f-string

Now, when we run add_logs("example.py", "logged_example.py") we get: 

def add(x, y):
    _ = x + y
    print(f"leave add({x!r}, {y!r}): {_!r}")
    return _
add(1, 2)

And when we run that, we get leave add(1, 2): 3, just like we wanted. To add the enter message, we can use the same format_call("enter") (that's why I made it accept an action and put the : {_!r} bit in an if-statement, sneaky me); we just have to act after the node (that is, the function definition) has been visited, so we have access to its fully transformed body and can insert our print statement at its beginning. Oh, and we should probably change the class name to something more appropriate: 

>>> class LogCalls(ast.NodeTransformer):
...     def visit_FunctionDef(self, node):
...         # … same as before
...         self.generic_visit(node)
...         print_ = ast.Expr(
...             value=ast.Call(
...                 func=ast.Name(id="print", ctx=ast.Load()),
...                 args=[self.format_call("enter")],
...                 keywords=[],
...             ),
...         )
...         node.body.insert(0, print_)
...         return node

And now, assuming we wrap add_logs in a script, we can do: 

$ python add_logs.py example.py logged_example.py
$ cat logged_example.py
def add(x, y):
    print(f'enter add({x!r}, {y!r})')
    _ = x + y
    print(f'leave add({x!r}, {y!r}): {_!r}')
    return _
add(1, 2)
$ python logged_example.py
enter add(1, 2)
leave add(1, 2): 3

Here I'd like to give a shout out to metap, a metaprogramming library (one of very few, really) that enlightened me as to the potential of this approach, which fits so snugly in our narrative. To borrow their terminology, what we've developed is a metaprogramming layer: that word that I was missing for a toolchain that rewrites code from a language unto itself.

Because that's what add_logs.py does, isn't it? It takes a source program (example.py), much like a compiler does, and transmutes it into something else (logged_example.py) – although here, like we said, it's not the lowered code a compiler produces, nor the translated one a transpiler does; it's, well, a metaprogrammed version thereof: analyzed and augmented by meta-code.

What else can we do, if we turn our add_logs.py into a more general-purpose layer, like metap is? Again, borrowing from their abstraction, a metaprogramming layer can:

  1. Augment programs: essentially what we did, rewriting some code into other code;

  2. Generate code: rewrite the program in such a way as to add additional functionality – and not one handling a common problem like logging, either, but custom logic defined by users;

  3. Inspect structure: expose an API, based on our newfound AST manipulation powers, that would let users ask questions about the code's structure.

The third point lets us do static analysis, like linters or type checkers do, which can help us assert the code adheres to certain standards or patterns – but it's the second point that I find particularly intruiging.

The Standard of Elegance

What would it mean, to have user-defined code generation? I mean, when users write code, we call it "programming"; when we talk of "code generation", it implies code is written automatically, in which case the methods by which users can customize it are not entirely clear. Sure, it can be context-aware, just like our enter and leave messages accounted for each function's particular arguments; and it can parameterizeable, letting users toggle options that generate code one way or another. But how can it be user-defined?

Luckily, we've already seen this paradigm in action, where constructs created by users are hoisted and integrated into a bigger whole: we called them plugins, and saw how they can accommodate all sorts of aspects, from web apps to tests to filtering and post-processing. In this case, we're talking about code rewriting, so let's consider what a plugin of this sort could mean. Before we do, though, I feel like a bit of history is in order. Consider this abstract code: 

>>> def f(x, y):
...     do_something(x)
...     do_something(y)
...     return x + y

Depending on what sort of something we want to do, programming languages already come equipped with ways of code composition: first and foremost the ability to call functions, of course. So, if we want to do something with no side effects, like printing x and y, we'd have no problem acheiving it with: 

>>> def f(x, y):
...     print(x)
...     print(y)
...     return x + y

If, however, we want to do something that modifies our values, we'd have to rely on another "trivial" feature: return values. To increment x and y, for example, we'd write: 

>>> def f(x, y):
...     x = inc(x)
...     y = inc(y)
...     return x + y

Obvious though it may seem, this wasn't always possible. While all programming languages always had the ability to jump – it's sort of a prerequisite for being Turing-complete – and could thus create reusable blocks of code, the interface that mediates the process evolved over time. Initially, to modify a value, we'd have to rewrite it in memory; to simulate this in Python, we need to wrap each value in a reference objects, often called pointers: 

>>> def f(x_obj, y_obj):
...     inc(x_obj)
...     inc(y_obj)
...     return x_obj.value + y_obj.value

So that inc can dereference its value and modify it in place: 

>>> def inc(obj):
...     obj.value = obj.value + 1

Furthermore, when return values were introduced, they were still limited: a function could return a single, simple value at most. This posed a challenge: what if a function might fail (that is, "return" an error) as well as produce a result? Suppose that instead of incrementing x and y, we'd want to get their reciprocals, 1/x and 1/y – this code: 

>>> def f(x, y):
...     x = recip(x)
...     y = recip(y)
...     return x + y

Would not do, since if either value is zero, an error indication is in order. In languages like C, we have to revert back to pointers, perserving the return value's limited capacity for status reports: 

>>> def f(x_obj, y_obj):
...     err = recip(x_obj)
...     if err:
...         return
...     err = recip(y_obj)
...     if err:
...         return
...     return x_obj.value + y_obj.value

Whereby: 

>>> def recip(obj):
...     if obj.value == 0:
...         return 1 # error
...     obj.value = 1 / obj.value
...     return 0 # success

In fact, if we're being pedantic, that's what our f should do as well: if either of its recips fails, it should fail as well. But this means it should return a status, too, which means its return value is occupied, which means that to communicate its result, we need a third, so-called output parameter: 

>>> def f(x_obj, y_obj, out_obj):
...     err = recip(x_obj)
...     if err:
...         return err
...     err = recip(y_obj)
...     if err:
...         return err
...     out_obj.value = x_obj.value + y_obj.value
...     return 0 # success

Pretty clumsy, no? And yet in C, such code is completely commonplace, even today. The point I'm trying to illustrate is just how profoundly the flow control capabilities of a language affect its clarity, flexbility, and overall design: as you can see, the inability to communicate errors by other means proliferates, shaping the entire codebase as a result. Of course, in Python this isn't the case – first of all, because we can return multiple values: 

>>> def f(x, y):
...     x, err = recip(x)
...     if err:
...         return None, err
...     y, err = recip(y)
...     if err:
...         return None, err
...     return x + y, 0

And second, because we have exceptions, so all recip needs to do is: 

>>> def recip(value):
...     if value == 0:
...         raise ZeroDivisionError()
...     return 1 / value

And we're back to: 

>>> def f(x, y):
...     x = recip(x)
...     y = recip(y)
...     return x + y

Because a ZeroDivisionError will tear through the call stack until it's caught; exceptions, in other words, are a flow control construct, and a pretty sophisticated one at that. My point, once again, is just how important such flow control is – and to what a degree it shapes our perception of how code can (and should) look. C programmers have no problem with error indications and output parameters; until multiple return values and exceptions were introduced, we didn't know what we were missing, so to speak – and even when they were, many criticized the latter for being too slow and too complex. In terms of performance, at least in realtime languages like C (well, C++), it's true that exceptions are considerably more expensive than returning an integer (although in high-level languages like Python and Javascript, it's completely immaterial); in terms of complexity, well, I just don't agree. It's always the same stubborn argument: "I am used to thinking about software a certain way, so a different paradigm must be necessarily worse." Google's C++ coding standards are a case in point: they initially forbade exception on these exact grounds, that raising and catching them makes control flow more confusing (read: "I'm not used to it") and less predictable (read: more flexible, if you learn to predict it). Over time, the excuse shifted to the irrefutable "our codebase doesn't use it, so adding it now would cause problems", but even they admit that exceptions are used by most other modern languages, and that their benefits probably outweigh their costs.

Personally, I don't care for branding language features as good or bad; these arguments always get very hypothetical ("what if someone does..."), and usually assume people are stupid ("bad programmers will..."). My attitude, as you can see from this book, is that people can learn and get smarter, and their creative experiments, even when they don't work out, carry latent and potentially explosive benefits. In terms of software design, I prefer to judge a codebase based on the clarity and novelty of its thought, not language. Learning new syntax is considerably easier and less exciting than learning new concepts, as those can fundamentally alter how we percieve and reason about a subject – just skim through the evolution of f again to see what I mean. That, I believe, is where the prize is – and why I see metaprogramming as such a treasure trove of untapped potential. The mere readability of code, its aesthetics – our current standard of elegance, if you will – is constantly changing as come up with new, more powerful paradigms; and the unspoken truth about it is that it can be taught, especially to people who are curious to learn.

User-Defined Code Generation

Which brings us back to user-defined code generation. With today's flow control, one of the things we still can't do is make do_something quit f. It can raise an error, sure, which then propagates – but if we want f to exit gracefully, we have to either catch it, or go back to status reports: 

>>> def f(x, y):
...     try:
...         do_something(x)
...     except Exception:
...         return
...     # … same for y
>>> # or...
>>> def f(x, y):
...     err = do_something(x)
...     if err:
...         return
...     # … same for y

What if there was a way to embed code – not call it in a different scope like we do with functions, but actually paste it in, so that if do_something is: 

>>> def do_something(value):
...     if value is None:
...         return

f would end up as: 

>>> def f(x, y):
...     if x is None:
...         return
...     if y is None:
...         return
...     return x + y

I guess you see where I'm going with this. Let's create a decorator, macro, which will be used to mark functions as embeddable: 

>>> @macro
... def return_if_none(value):
...     if value is None:
...         return

In practice, though, it will make them uncallable: 

>>> def macro(function):
...     def wrapper(*args, **kwargs):
...         raise RuntimeError("metaprogramming layer must run first")
...     return wrapper

Because, like the error states, for this to work we have to run it through a metaprogramming layer first, giving it a chance to actually embed our macro. Since it works on source code, let's put everything in appropriate files; in example.py we have: 

from meta import macro

@macro
def return_if_none(value):
    if value is None:
        return
        
def f(x, y):
    return_if_none(x)
    return_if_none(y)
    return x + y
    
print(f(1, 2))
print(f(1, None))

That is, some code – f and a couple of its invocation – and the return_if_none macro, which can be thought of as bona fide user-defined code-to-be-generated. In meta.py we'll put our silly macro decorator, as well as the transformer that will take care of embedding it and the API to run it: 

import ast, sys

def macro(function):
    # … same as before
    
class Embedder(ast.NodeTransformer):
    # … to be continued
    
def embed(source, target):
    code = open(source).read()
    tree = ast.parse(code)
    tree = Embedder().visit(tree)
    ast.fix_missing_locations(tree)
    open(target, "w").write(ast.unparse(tree))
    
if __name__ == "__main__":
    embed(sys.argv[1], sys.argv[2])

Now for the transformer itself: the first thing we need to detect is all the function definitions, so as to see if they're decorated with @macro – and if so, keep their node for later, to be embedded on demand: 

class Embedder(ast.NodeTransformer):
    def __init__(self):
        self.macros = {}
    def visit_FunctionDef(self, node):
        for decorator in node.decorator_list:
            if isinstance(decorator, ast.Name) and decorator.id == "macro":
                self.macros[node.name] = node
                return None # remove definition from final code
        self.generic_visit(node)
        return node

Next, we need to find all the call sites, and if they're invoking a macro, replace them with its actual content. Doing so in visit_Call will not do, since those are expressions that can't be replaced with a statement; so instead, we'll target the expression-statement (Expr) nodes that wrap them: 

class Embedder(ast.NodeTransformer):
    # … same as before
    def visit_Expr(self, node):
        if isinstance(node.value, ast.Call):
            name = node.value.func.id
            if name in self.macros:
                return self._embed(self.macros[name], node)
        return node

The way we do the actual embedding is a bit tricky. We can't simply return the macro body as-is, because then we'll end up with: 

def f(x, y):
    if value is None:
        return None
    # …

See the problem? What we need is a mapping between the macro's parameters and the arguments passed into the call – so we traverse the body and swap each such name with its corresponding value: 

import copy

class Embedder(ast.NodeTransformer):
    # … same as before
    def _embed(self, macro, call):
        args = {}
        for param, arg in zip(macro.args.args, call.value.args):
            args[param.arg] = arg
        embed = Substituter(args).visit(copy.deepcopy(macro))
        return embed.body

Admittedly, it's pretty confusing, but each macro.args.args's .arg refers to a parameter name in the macro's signature, while each of call.value.args is the argument that's been passed to the call. And since visitation changes nodes in place, note that we need to create a deep copy of the macro each time we embed it. As for the substitution itself, there's not much to it: 

class Substituter(ast.NodeTransformer):
    def __init__(self, args):
        self.args = args
    def visit_Name(self, node):
        return self.args[node.id] if node.id in self.args else node

With that, we're done: 

$ python meta.py example.py embedded_example.py
$ cat embedded_example.py
from meta import macro # should've probably removed this, too

def f(x, y):
    if x is None:
        return
    if y is None:
        return
    return x + y
print(f(1, 2))
print(f(1, None))
$ python embedded_example.py
3
None

And I know, I know; it's quite a lot of work, where we could have simply written if x is None: return by ourselves. But that's not the point – just an example that might reasonably fit in a book. That's why I went on that weird tangent about return values and exceptions earlier: what we've implemented here is a new flow control construct, and that, as we've seen, is no small feat. Instead of jumping into functions, we can embed them directly, unlocking possibilities that simply did not exist in the language before. What design patterns and programming styles it might lead to is hard to predict – but as far as our adventure goes, it's a pretty cool item to cross off our bucket list.