Assertion rewriting in Pytest part 3: The AST

Part 1 and part 2 looked at why a test framework might need to handle method calls in an unusual way, and how we might go about implementing this by the simplest possible means.

The approach in part 2 has a couple of limitations, but one is that we’re using eval, which has a very limited interface. If we have something like:

result = eval("len(some_list) > 10")

then eval lets us find out whether the condition succeeds or fails, which is useful. But ideally we want to go further, and figure out how exactly the condition failed. We know that len(some_list) was less than or equal to 10, but exactly what was it? Maybe it would be nice to see the elements of some_list printed out to aid debugging. eval doesn’t give us any of this.

The problem is that eval is actually doing too much at once. There are broadly three steps:

  • Parsing: What is the structure of len(some_list) > 10? This stage converts a simple string of characters into a richer structure whose shape encodes everything we need to know to execute the code, and nothing more.
  • Compiling: Convert this into a form that the computer can execute.
  • Execution: Find the outcome from executing the code with the specific input values in question.

If we can split these apart, then it’s much easier to write code that manipulates the behaviour of the piece we are interested in.

Let’s look at the first stage for now.

I glibly talked about the structure having “everything we need, and nothing more”, but that’s not really a practical starting point. What should such a structure actually look like?

Without digressing too much, here’s an example of the sort of structure that’s used in Python:

 

If you look at this from the top down, you’ll first see a less-than (<) operator that has two arguments. The right hand argument is a simple number, but the left-hand argument is more complex: it’s made up of a function (len) applied to a further argument (some_list). Some points about this structure:

Firstly, it’s not necessarily obvious from the trivial example here, but this is a tree structure. As you look further down it can split apart into different branches, but the branches never join up with each other to form loops. This makes the structure much easier to work on at the cost of limiting what we can express with it; it turns out that this particular trade-off is a good one.

Secondly, the value of a node in the tree is only affected by the nodes below it, not all the other nodes in the tree. For the len node, the function has the same value whether it’s appearing in len(some_list) < 10 or len(some_list) > 6. This corresponds to what we expect, and is a good sign that this structure is going to make our life easier than if we used a different structure.

This sort of structure is called an Abstract Syntax Tree or AST and as I say, most languages use something like this. However, Python is relatively unusual in the degree to which it exposes the AST to you as a programmer. You can just import the ast module and start playing around:

>>> import ast
>>> tree = ast.parse('len(some_list) > 10', mode='eval')
>>> ast.dump(tree)
"Expression(body=Compare(left=Call(func=Name(id='len', ctx=Load()),
args=[Name(id='some_list', ctx=Load())], keywords=[]), ops=[Gt()],
comparators=[Num(n=10)]))"

If I rearrange that dump a little bit (and cut out some details), we can see:

Expression(
    body=Compare(
        left=Call(
            func=Name(id='len', ...),
            args=[
                Name(id='some_list', ...)
            ],
            ....
        ),
        ops=[Gt()],
        comparators=[
            Num(n=10)
        ]
    )
)

It doesn’t take too much imagination to see that the diagram above is a rough approximation to this.

So we’re peering into the first stage of executing the code; we’re actually seeing inside the workings of exec. But that’s not all. This is real working Python code, and we can prove this by carrying on with the compile and execute stages:

>>> bytecode = compile(tree, '<string>', mode='eval')
>>> eval(
        bytecode,
        {},
        {
            'some_list': [1, 2, 3]
        }
    )
False

We just did the compile and execute (second and third) stages ourselves, and got the same output as if we had let Python do it for us inside of exec.

It’s worth reiterating: this isn’t just a curiosity. This is a practical tool that you can use in real Python code, and you can do this with just the tools in the standard library. I’ll look at how this might be used in practice in the next part.

Leave a Reply

Your email address will not be published. Required fields are marked *