DEV Community: Tiago Antao

Automatic code changing in Python with the ast module

Tiago Antao — Sun, 05 Dec 2021 23:57:40 +0000

In this article we are going to explore Python Abstract Syntax Trees. Our example will use the Python ast module to convert legacy Python 2 to Python 3.

One of the changes that Python 3 brings versus Python 2 is that, if you redefine the __eq__ method of an object then your __hash__ method automatically becomes None. This means that your class is rendered unhashable (and thus unusable in many settings). This makes sense as the
hash function is closely related to object equality. But, if you have Python 2 code that is nonetheless working and you are planning to move to Python 3, then automatically adding an __hash__ function to recover object hashbility might be a pragmatically acceptable solution.

This not-so-perfect example sets the stage for a mini-tutorial on traversing a Python source tree, finding classes that have new __eq__ methods but lack a new __hash__ method and then patching the class to have one.
This is, thus, mostly an excuse for a brief exploration on the ast module of Python...

The ast module in Python

The ast module allows you to access the Abstract Syntax Tree of a Python program. This can be useful for program analysis and program transformation, among other things. Lets start with the AST representation of a Hello World program:

# Typical hello world example
print('Hello World')

Getting the AST tree for this is trivial:

import ast
tree = ast.parse(open('hello.py').read())

Now if you decide to explore the tree object, it might seem daunting at first, but don't let appearances fool you, it is actually quite simple. The type of tree will
be ast.Module. If you ask, by doing tree.fields, for its children
(as you ask can to all AST nodes) you will see that a Module will have a
body (i.e. tree.body). The body attribute will have, you guessed it, the body of the file:

print(tree.body)
<_ast.expr object at 0x7f5731c59bb0>

The body of a module is composed of a list of Statements. In our case
we have a single Node: an Expr (Expression is another type of node -
do not confuse them). Notice that the AST representation will not include the comments, these are "lost".

OK, back to our Call. What is in a call? Well, you call a function with arguments so a call is a function name plus a set of arguments:

print(tree.body[0].value._fields)
('func', 'args', 'keywords', 'starargs', 'kwargs')

You will notice the func plus a lot of stuff for arguments. This is because, as you know, Python has quite a rich way of passing arguments (positional arguments, named arguments, ...). For now lets concentrate
on func and args only. func is a Name with an attribute called id
(the function name). args is a list with a single element, a string:

>>> print(tree.body[0].value.func.id)
print
>>> print(tree.body[0].value.args[0].s)
Hello World

This is actually quite simple, here is a graphical representation:

The best way to find all types of nodes is to check the abstract grammar documentation.

Second attempt (a function definition)

Lets have a look at the AST for a function definition:

def print_hello(who):
    print('Hello %s' % who)

If you get a tree for this code, you will notice that the body is still composed of a single statement:

>print(tree.body)
<_ast.functiondef object at 0x7f5731c59bb0>

At the top level, there is only a single statement, the print function
call belongs to the function definition proper. A function definition
has name, args, body, decorator~list and returns. Name is a string
print_hello, no more indirections here, easy. args cannot be very
simple because it has to accommodate the fairly rich Python syntax for
argument specification, so its fields are args, vararg,
kwonlyargs, kw_defaults, kwarg, defaults. In our case we just
have a simple positional argument, so we can find it on args.args:

print(tree.body[0].args.args)
<_ast.arg object at 0x7f5731c59bb0>

The args object has a arg field (tree.body[0].args.args[0].arg starts to sound a bit ludicrous, but such is life), which is a string
(who). Now, the function body can be found in the body field of the
function definition:

print(tree.body[0].body)
<_ast.expr object at>

The body is as discussed above for the print call.

To finalize this version, I just want to present the sub-tree inside the
print("Hello %s" % who) - Notice the BinOp for the % operator:

OK, before we go to the final version, one very last thing:

# Lets look at two extra properties of the print line...
print(tree.body[0].body[0].lineno, tree.body[0].body[0].col_offset)
2 4

Yes, you can retrieve the line number and the column of a statement...

Third attempt

class Hello:
    def __init__(self, who):
        self.who = who

    def print_hello(self):
        print('Hello %s' % self.who)

It should not come as a shock that the module only has a single statement:

<_ast.classdef object at>

And yes, as you might expect by now, the ClassDef object has a name
and a body attribute (things start making sense and are consistent).
Indeed most things are as expected: The class has two statements (two
FuncDefs). There are only a couple of conceptually new things, and
these are visible in the line:

self.who = who

Here we have two new features: the assignment = and the compound name
self.who.

Notice that the list of potential targets is a list, this is because you
can do things like:

x, y = 1, 2

There is much more than can be said. Processing ASTs requires a
recursive mindset, indeed if you are not used to think recursively I
suspect that might be your biggest hurdle in doing things with the AST
module. And with Python things can get very nested indeed (nested
functions, nested classes, lambda functions, ...)

A complete solution

This code was run a few years ago. While it is for Python 2, the general thought process is the same with Python 3

OK, lets switch gears completely and apply this to a concrete case. The
Abjad project helps composers build up
complex pieces of music notation in an iterative and incremental way. It
is currently Python 2 only, and I am volunteering some of my time to
help it support Python 3. This is highly-OO, highly well documented and
with a massive load of test cases (kudos to the main developers for
this). Some of the classes do have __eq__ methods defined, but lack
__hash__ methods, so a pragmatic (though not totally rigorous
solution) is to add the __hash__ methods required by Python 3.

A caveat here: the processing has to be done in Python 2, so the code
below is Python 2. The idea is to generate Python 2 code with hashes
that can be automatically translated by 2to3 (no need for monkey
patching after 2to3)

First thing, we have to traverse all the files:

def traverse_dir(my_dir):
    content = os.listdir(my_dir)
    for element in content:
        if os.path.isdir(my_dir + os.sep + element):
            traverse_dir(my_dir + os.sep + element)
        elif os.path.isfile(my_dir + os.sep + element) and element.endswith('.py'):
            process_file(my_dir + os.sep + element)</pre>

Nothing special where, just plain traverse of a directory structure. Now
we want to find all classes that have an __eq__ method, but not an
__hash__ method:

def get_classes(tree):
    # Will not work for nested classes
    my_classes = []
    for expr in tree.body:
        if type(expr) == ast.ClassDef:
            my_classes.append(expr)
    return my_classes</code></pre>

The function above will return all classes from a list of statements
(typically a module body).

def get_class_methods(tree):
    my_methods = []
    for expr in tree.body:
        if type(expr) == ast.FunctionDef:
            my_methods.append(expr)
    return my_methods

The function above will return all function definitions from a
ClassDef object.

def process_file(my_file):
    shutil.copyfile(my_file, my_file + '.bak')
    tree = ast.parse(open(my_file).read())
    my_classes = get_classes(tree)
    patches = {}
    for my_class in my_classes:
        methods = get_class_methods(my_class)
        has_eq = '__eq__' in [method.name for method in methods]
        has_hash = '__hash__' in [method.name for method in methods]
        if has_eq and not has_hash:
            lineno = compute_patch(methods)
            patches[lineno] = my_class.name
    patch(my_file, patches)

This is the main loop applied to each file: We get all the available
classes; for each class we get all available methods and if there is an
__eq__ method with no __hash__ method then a patch is computed and
then applied. The first thing that we need is to know where to patch the
code:

def compute_patch(methods):
    names = [method.name for method in methods]
    names.append('__hash__')
    try:
        names.remove('__init__')
    except ValueError:
        pass
    names.sort()
    try:
        method_after = names[names.index('__hash__') + 1]
    except IndexError:
        method_after = names[-2]
    for method in methods:
        if method.name == method_after:
            return method.lineno

The main point here is to decide to which line to apply the patch. It
has of course to be inside the class, but we want to be slightly more
elegant than that: We want to put the method in lexicographical order
with all the others (so, __hash__ would go between __eq__ and
__init__). Now we can patch:

def patch(my_file, patches):
    f = open(my_file + '.bak')
    w = open(my_file, 'w')
    lineno = 0
    for l in f:
        lineno += 1
        if lineno in patches:
            w.write("""    def __hash__(self):
        r'''Hashes my class.

        Required to be explicitely re-defined on Python 3 if __eq__ changes.

        Returns integer.
        '''
        return super(%s, self).__hash__()

""" % patches[lineno])
        w.write(l)

Final notes

This could be done in many other ways. Indeed this is not even the standard way (that would be coding a 2to3 fixer). The practical solution is not general (for instance, it does not support nested classes). Also,
it has problems with comments (as we lose them in the AST). But, for the
practical purpose (patching abjad) it was good enough (You can see the result here.

For further reading, you might want to have a look at this stack overflow question.

Advanced Array access in Cython using a practical example

Tiago Antao — Wed, 09 Jun 2021 13:52:02 +0000

This is an Excerpt from https://www.manning.com/books/fast-python-for-data-science available from Manning.

Python is slow. The standard implementation is slow and the language dynamic features pay a performance toll. Many Python libraries are performant precisely because they are partially implemented in lower level languages making available efficient data processing algorithms. But sometimes we will need to implement our own high-performance algorithms in something faster than Python. In this article we are going to consider Cython, a superset of Python that is converted to C and is substantially more performant than Python.

We are going to discuss Cython via a concrete example: a color version of Conway’s Game of Life. Conway’s Game of Life is a zero player game which evolves automatically from its initial state — devising interesting initial states is part of the fun. The state of the game is composed of a grid of arbitrary size and each cell can have two states: alive or dead. As time advances each cell will change its state according to the following rules:

Any live cell with 2 or 3 neighbors survives
A dead cell with 3 neighbors becomes live
All other cells die or stay dead

The world wraps around in the sense that the left-most column will look at the right-most column to compute neighbors (and vice-versa). The same applies for the top and bottom rows.
There are 3 examples shown over time in figure 1: the first example is a dash that will be eternally changing directions from vertical to horizontal; the second is a box that is stable and the third dies off completely.

We will be using an extension called QuadLife where each live cell can have 4 different states: Red, Green, Blue and Yellow. We prefer this extension simply because it looks cooler. It includes two new rules for newborns:

If there is a color that is present in the majority of neighbors, it becomes the color of the new cell
If all 3 live neighbors are different then the new cell takes the remaining color of the four possible

A Cython implementation of QuadLife

All the code referred here can be found in https://github.com/tiagoantao/python-performance/tree/master/05-cython/sec6-quadlife

When we say “Cython implementation”, we actually have two parts: the Cython part – which should cover the computational expensive inner loop and all the rest which is done in Python. As such our implementation will include two parts: A calling Python code with the computational expensive part implemented in Cython.

This example is of intermediate complexity from a Cython perspective: there are plenty of introductory tutorials on the Web and we want to add some value above the basic. That being said it should be self-contained.

The Python part has to refer to be able to link Cython. There are several ways to do that, here we will use pyximport which is a simple way to do that – as we don’t want to spend much time in this article discussing linking.

import sys

import numpy as np
import pyximport
pyximport.install(   (1)
    language_level=3,
    setup_args={
        'include_dirs': np.get_include()})

import cquadlife as quadlife

SIZE_X = int(sys.argv[1])   (2)
SIZE_Y = int(sys.argv[2])
GENERATIONS = int(sys.argv[3])

world = quadlife.create_random_world(SIZE_Y, SIZE_X) (3)
for i in range(GENERATIONS):     (4)
    world = quadlife.live(world)

Notes above:

We setup ~pyximport~ to include NumPy
We read parameters from the command line
We use a function — defined later — to create a random world
We apply the Quadlife algorithm for the number of generations specified by the user

We call this script by passing the desired X and Y resolution along with the number of generations. The script doesn’t output anything for now, it just runs the game — later we will do some cool stuff with the results.

To start we will generate a random world with create_random_world — later we will consider better looking alternatives — which is good enough for testing. We will use a NumPy array of dimension SIZE_Y, SIZE_X as specified by the user. It will be filled with random values between 0 and 4. 0 represents a dead cell. We then run the simulation function called live for GENERATIONS generations: the first call gets the random world and its output is then fed to itself in succession.

We now consider our Cython code. Creating a random world doesn’t need to be really optimized as it is called only once at the beginning:

#cython: language_level=3
import numpy as np

cimport cython
cimport numpy as cnp

def create_random_world(y, x):
    cdef cnp.ndarray [cnp.uint8_t, ndim=2] world = np.random.randint(0, 5, (y, x), np.uint8)
    return world

Notice that we are typing the variable world: this is a Cython extension that allows to generate C code that is more specific to our data and hence faster.
Now the fun starts. First off, we want to make sure we our inner loop can be GIL-free. For that we create a Cython top-level live function that mostly converts NumPy arrays into memoryviews.

def live(cnp.ndarray[cnp.uint8_t, ndim=2] old_world):
    cdef int size_y = old_world.shape[0]
    cdef int size_x = old_world.shape[1]
    cdef cnp.ndarray[cnp.uint8_t, ndim=2] extended_world = np.empty((size_y + 2, size_x + 2), dtype=np.uint8)
    cdef cnp.ndarray[cnp.uint8_t, ndim=2] new_world = np.empty((size_y, size_x), np.uint8)
    cdef cnp.ndarray[cnp.uint8_t, ndim=1] states = np.empty((5,), np.uint8)

    live_core(old_world, extended_world, new_world, states)
    return new_world

The actual conversion to memoryviews will be forced by the live_core function signature — see below — but we still need a layer that is able to convert Python objects into potentially GIL-free representations (objects that do not require the Python machinery and thus bring us to a situation where we can release the GIL). old_world is the input world, new_world will have the output. extended_world and states are live_core internal variables that we will pre-allocate here. Before we present the core algorithm in live_core lets discuss how we are going to algorithmically optimize part of it.

In the game of life, the extremes of the board are connected, for example the left-most column cells will "look" at the states of the right-most cells to compute its new state. In order to avoid a lot of corner case testing — which would add a lot of if statements and hence time to compute — we will implement a temporary extended world — in the above mentioned variable extended_world of dimensions (y+2, x+2). The extended boundaries copy what happens on the other side, as in figure 2.

The purpose of this algorithm is that it allows a slightly more efficient approach when computing the new board: we do not need to if for boundary conditions. This is done at the cost of memory — we now need to store a new and bigger version of the board. Making these kinds of tradeoffs — memory vs computation — is our bread and butter across high performance computing problems. It’s difficult to come up with generalized guidelines to make a decision on the tradeoff: It will depend on the computational and memory cost of a specific algorithm and the resources that you have available.

Here is the code that implements this extended world — notice how the code has no boundary tests thus reducing the computation time of using if:

@cython.boundscheck(False)  (1)
@cython.nonecheck(False)
@cython.wraparound(False)
cdef void get_extended_world(   (2)
        cnp.uint8_t[:,:] world,
        cnp.uint8_t[:,:] extended_world): (3)
    cdef int y = world.shape[0]
    cdef int x = world.shape[1]
    extended_world[1:y+1, 1:x+1] = world (4)

    extended_world[0, 1:x+1] = world[y-1, :]  # top
    extended_world[y+1, 1:x+1] = world[0, :]  # bottom
    extended_world[1:y+1, 0] = world[:, x-1]  # left
    extended_world[1:y+1, x+1] = world[:, 0]  # right

    extended_world[0, 0] = world[y-1, x-1]   # top left
    extended_world[0, x+1] = world[y-1, 0]   # top right
    extended_world[y+1, 0] = world[0, x-1]   # bottom left
    extended_world[y+1, x+1] = world[0, 0]   # bottom right

Notes:

Use deactivate bounds, None and wrap checks
We use cdef to be able to avoid the GIL
We type everything on the function signature
The copy of world into the middle of extended_world can be potentially expensive

The copy of world into the middle of extended_world can potentially pay a steep computational — along with memory — price, but the computational part might be compensated by the easier core algorithm. But at least for pedagogical purposes, it makes the core algorithm substantially simpler which is important for learning.

You might notice that many lines on the function above look like they could be written in a more expedient notation. For example, perhaps:

extended_world[1:y+1, 1:x+1] = world

Could maybe be written as:

extended_world[1:-1, 1:-1] = world

It turns out, however, that we cannot do these kinds of rewrites because we deactivated wraparound in order to avoid paying the price of the generated C code to spend time on wrap verification. Deactivating wraparound means we cannot use negative indexes. The tradeoff is worth it, though: wrap around requires the CPython machinery, which slows things down, so our implementation without it is substantially faster and we need to deactivate it to release the GIL as the wrap-around could would use the Python machinery.

NOTE: Not doing wrap around or bounds checking might result
in Segmentation faults of your code. In you see those
errors make sure to deactivate the decorator during
development. You need to make sure your code is robust
enough to tolerate removing this and the other checks.

We also use other optimizations: a cdef (a function that only has a C-interface and cannot thus be called from Python), complete typing of parameters and variables and use of memoryviews instead of NumPy arrays. If you use cython -a cquadlife.pyx (-a generates a HTML with code interations with C and the CPython machinery) you will see that there is no yellow — i.e. Python interaction lines — in the code above.

The main implementation that changes the state makes use of the extended world. The code below implements the QuadLife game rules. Because it is quite long we will carefully annotate it, including issues that we might have alluded upon in the past.

@cython.boundscheck(False)    (1)
@cython.nonecheck(False)
@cython.wraparound(False)
cdef void live_core(  (2)
    cnp.uint8_t[:,:] old_world,  (3)
    cnp.uint8_t[:,:] extended_world,
    cnp.uint8_t[:,:] new_world,
    cnp.uint8_t[:] states):   (4)
    cdef cnp.uint16_t x, y, i   (5)
    cdef cnp.uint8_t num_alive, max_represented
    cdef int size_y = old_world.shape[0]
    cdef int size_x = old_world.shape[1]
    get_extended_world(old_world, extended_world)  (6)

    for x in range(size_x):
        for y in range(size_y):
            for i in range(5):
                states[i] = 0
            for i in range(3):
                states[extended_world[y, x + i]] += 1
                states[extended_world[y + 2, x + i]] += 1
            states[extended_world[y + 1, x]] += 1
            states[extended_world[y + 1, x + 2]] += 1

            num_alive = states[1] + states[2] + states[3] + states[4] (7)
            if num_alive < 2 or num_alive > 3:
               # Too few or too many neighbors
                new_world[y, x] = 0
            elif old_world[y, x] != 0:
                # Stays alive
                new_world[y, x] = old_world[y, x]
            elif num_alive == 3:  # Will be born
                max_represented = max(states[1], max(states[2], max(states[3], states[4]))) (8)
                if max_represented > 1:
                    # majority rule for color
                    for i in range(1, 5):
                        if states[i] == max_represented: (9)
                            new_world[y, x] = i
                            break
                else:
                    # diversity - use whichever color doesn't exist
                    for i in range(1, 5):
                        if states[i] == 0: (10)
                            new_world[y, x] = i
                            break
            else:
                new_world[y, x] = 0  # stays dead

Notes:

We deactivate a lot of checking machinery: bounds and None checking along with wrap-around
We use a cdef to avoid passing standard Python objects. We also declare the return type as void which is C for nothing at all
We type all parameters
Some internal variables — states and extended_world — were actually allocated outside and we use the memory made available
We type all local variables
Everything is pre-allocated when we call get_extended_world`
Implements sum(states[:1])
Implements max(states[:1])
Implements states[1:].index(max_represented)
Implements states[1:].index(0)

This function is complicated but many of the techniques used here were actually introduced before — they are actually combined in a more realistic example. So read the code carefully and you will understand it all.

You might see a few oddities in the code, namely the replacement of sum and index by less declarative versions. We do this because sum and index would make use of the CPython machinery, and we want to avoid that. A similar argument is valid for the max function, but in that case there is a replacement version that is able to compare values at a time. When you are using general functions, you might want to profile them and maybe replace them with optimized non-general functions.

NOTE: Because the game of life produces nice evolving visualizations, we are going to make a simple GUI. We will be using the Python built-in tkinter module for the GUI along with the external library Pillow for image manipulation.
Our implementation is now complete but we would like to gauge the performance gain that we got from our code

Basic performance analysis

You will be able to find a native Python version on the repository, we will use this version to make some basic comparisons with the Cython one. Running the Python version for a resolution of 1000x1000 for 200 generations on my computer takes slightly less than 1000 seconds, i.e. a bit less than 17 minutes. The Cython code takes 2.5 seconds.

Our implementation is memory intensive. Be very careful if you test it with large resolutions.

Let’s take a lighter approach and generate a cool video from the game and make some considerations about computational complexity—especially with regards about the role of theory to help us writing more efficient programs.

A spacewar example using Quadlife

In the repository you will find code that will generate a video using a starting state that includes "spaceships" and "defenses" (The patterns come from http://www.conwaylife.com). The code itself is not very relevant for optimization purposes so we won’t discuss it here. If you want to replicate it you will need Python’s Pillow library for image processing and ffmpeg to generate videos.

Here is the starting state, color inverted:

And here is the video link

This code will run the game of life on the spaceship model for 400 generations on a map of 400x250 and takes less than a second. For an HD resolution of 1920x1080 with the same 400 generations it would take around 11 seconds, 800 frames would take 22 seconds. 400 frames at the 4K resolution of 3840x2160 would take 48 seconds. At 40 frames per second a 4K game of 90 minutes would take roughly 196 minutes to generate. The silver lining is that the same video would take 54 days to generate in a pure Python solution.

This is an Excerpt from https://www.manning.com/books/fast-python-for-data-science available from Manning.