“Everyone knows that debugging is twice as hard as writing a program in the first place. So if you’re as clever as you can be when you write it, how will you ever debug it?” Brian W. Kernighan

In previous post I described the way to use Python AST for static analysis.

In this post I will show you how to add guard expressions support to Python with (relatively) simple Abstract Syntax Tree modification.

Here is syntax we want to support:

class ProfileHandler:
    @when(request.user.is_authenticated() and request.user.id == profile_id)
    def get(self, request, profile_id):
        # user is profile owner
    @when(request.user.is_authenticated())
    def get(self, request, profile_id):
        # user is authenticated, but not profile owner
    def get(self, request, _):
        # default case - anonymous user

It may seem Django-like, but thats because I wanted to show you, that we are going to create something really useful or at least illustrative.

Most articles I read about AST modification use basic examples - for example ast.Name transformation into function call or attribute fetching. While it may be interesting at first, those examples doesn’t explain you how to use AST modification in real world.

In this post I want to show that we can go much further with AST modification - we can create new constructions of Python language and even create our own, brand new and of course much better version of Python :-)

To be honest it’s really tough to create brand-new language based on Python syntax - language design is obscure, dark voodoo for mere mortals. But still, I believe that it’s possible to use this voodoo in small amounts to create better things - pytest is a great example of AST modification done right. The whole idea of fixtures tree is possible only thanks to AST modification.

In this post we are not going to create production-ready, robust and compatible version of Python. The ultimate goal of this post is to show you the way to augment Python to make you life easier (or turn into hell - it’s up to you).

High-level overview of required modifications

The only way I know to simulate guard expressions in Python is to use the if/elif/.../else construct.

Our goal is to glue several cases of one function together in correct order:

We also want to support cases when there is no default function defined.

How to

I tried to structure source code so that you can read it from top to bottom.

import ast
class Visitor(ast.NodeVisitor):
    def visit_ClassDef(self, node):
        # create empty dict for each class we process
        self.guarded_functions = {}
        # proceed to children elements of class
        # calls visit_FunctionDef for every method of this class
        # so after return from it we have populated dictionary of guarded
        # functions (thanks to our implementation of visit_FunctionDef)
        self.generic_visit(node)
        for name, cases in self.guarded_functions.items():
            # filter out methods with the same name as guarded function
            node.body = [el for el in node.body
                         if getattr(el, 'name', '') != name]
            # create generic function using first function case AST Node as base
            new_func = cases[0]
            # change original body with reassembled one
            new_func.body = self._build(cases[0], cases[1:])
            # remove decorators from function. @when is not real decorator,
            # so letting it stay will cause NameError
            new_func.decorator_list = []
            # add newly crafted AST Node to class node
            node.body.append(new_func)
    def visit_FunctionDef(self, node):
        guard = None
        for decorator in node.decorator_list:
            # find methods with @when decorator and fetch guard expression
            if isinstance(decorator, ast.Call) and decorator.func.id == 'when':
                guard = decorator.args[0]
                break
        name = node.name
        if guard:
            # if the FunctionDef was decorated with @when, add this node
            # to dict of guarded functions and save guard expression
            # in FunctionDef itself for further processing
            self.guarded_functions.setdefault(name, [])
            self.guarded_functions[name].append(node)
            node._guard_expression = guard
        elif name in self.guarded_functions:
            # if current function wasn't decorated with @when, but
            # the name is registered in guarded functions than it must be
            # default case of the function
            self.guarded_functions[name].append(node)
    def _build(self, func, remaining):
        # build new FunctionDef body recursively
        if remaining:
            # there are other cases, so we need to check current function's
            # guard expression and add recursively other cases to orelse
            return [ast.If(
                test=func._guard_expression,
                body=func.body,
                orelse=self._build(remaining[0], remaining[1:])
            )]
        else:
            if hasattr(func, '_guard_expression'):
                # this case is last one and there is _guard_expression,
                # so we need to create default case by ourself
                return [ast.If(
                    test=func._guard_expression,
                    body=func.body,
                    orelse=[ast.Raise(
                        lineno=func.lineno,
                        col_offset=func.col_offset,
                        exc=ast.Call(
                            ast.Name('TypeError', ctx=ast.Load()),
                            args=[ast.Str(s='Default case not found')],
                            keywords=[]),
                        cause=None)]
                )]
            else:
                # this case is the last one and there is no _guard_expression,
                # so we use it as default case
                return func.body
def compile_module(filename):
    with open(filename) as f:
        code = f.read()
    tree = ast.parse(code)
    visitor = Visitor()
    visitor.visit(tree)
    ast.fix_missing_locations(tree)
    return compile(tree, filename, 'exec')

Wrapping up

I want to rephrase Brian W. Kernighan: “AST modification is a razor-sharp tool, with which one can create an elegant and efficient program or a bloody mess”

Unfortunately, Python does not offer developer-friendly tools to modify itself, it’s AST or it’s lexicon. While I was writing this post, I saw many different, sometimes bizarre compilation errors. Sad thing about Python compilator is that it assumes that everything is going to be okay with syntax and it will not tell you where exactly you missed something. In this example, AST modification is relatively small, but I still had to read it carefuly to find bugs, because compiler refused to tell anything about it (special thanks to author of astpretty - it saved my life many times)

Despite the fact that presented “augmented” Python can’t be imported as an ordinary module, there are ways to change module importing machinery to make such modules look the same as those from standart Python. It is another dark, mysterious part of Python language and I hope that some day I’ll write about how to create custom modules using custom Finder and Loader classes injected in Python importing machinery.