Speeding up Python code with Codon

October 2, 2023

Last year I participated in the Advent of Code. It was a ton of fun and I learned some useful things along the way, even while being in my comfort zone and using Python. (You can check out my solutions here if you like!)

While improving some of my solutions on the way and using better algorithms I still stumbled on the limitation of Python in comparison to a compiled language. Especially in the challenges from day 15 on my code would sometimes run for some seconds, or even minutes to complete.

I then had a look around to see if there was a possibility to improve the runtime of my code, and indeed, there was! Besides trying nuitka (which didn't bring any improvements), I discovered Codon!

Although "not a drop-in replacement", it aims to be as close as possible to the Python syntax, and smaller scripts can run out of the box. You'll not be able to do things like

my_list = ['a', 5]

but it turns out I literally never use that. The other thing it needs is type hints (introduced with PEP 484), which I already use extensively (and makes working with Python code much easier), like here:

class Item:
    idx: int
    n: int
    left: Optional[Item]
    right: Optional[Item]

    def __init__(self, idx: int, n: int):
        self.idx = idx
        self.n = n
        self.left = None
        self.right = None

    def set_neighbors(self, left: Item, right: Item):
        self.left = left
        self.right = right

Notice the type definition of the instance variables idx, n etc. where you would normally find the class variables.

To run Codon, follow the readme on their Github page to install it, and after you can run it as follows:

codon run -release my_script.py

Timing

So let's have a look at some runtime comparisons! I took days 15-20 from AoC and day 23 as the longest-running code I had. These are the results (ordered by Speedup):

Day LoC Python 3.11.1 Codon 0.16.2 Speedup
19 152 332.88s 126.19s 2.6x
16 125 152.89s 42.52s 3.6x
17 175 0.34s 0.03s 13.2x
20 96 10.59s 0.61s 17.3x
15 108 5.14s 0.17s 30x
18 60 11.85s 0.38s 31.2x
23 123 5875.35s 87.3s 67.3x

The speedup was between 2.6x and 67.3x, with two scripts being below 10x, two being between 10-20x and 3 being >30x. I'm sure it heavily depends on the features used (and what the compiler can optimize), but from a "end-user" developer perspective it's good to know that the resulting speedup can vary heavily.

Gotchas

Besides the mixed collections there are other features of Python which are not trivial to run as-is. Some examples I stumbled upon when trying to run pure Python code:

  • import typing not working => reported & resolved (Github issue)
  • min() with key not working => reported (Github issue)
  • Imports: If you want to use other libraries (written in C or Python), it gets a bit complicated, but it's possible. You can also use the @python decorator, but wont be able to speed up that part of the code.
  • Class properties: Those have to be defined in from the start before being able to use them with self.my_property (see Item-class above).
  • Class names: To reference a class in itself you'll have to use it without '' (like neighbor: Node), where in pure Python it would be neighbor: 'Node'.
  • Function order: You can't reference a function that is defined later in the script (whereas in Python you can)
  • input() for getting user input does not work.
  • Equality, hash-code (__ne__, __hash__) and other "magic methods" must sometimes be explicitly implemented for custom classes.
  • Explicit Optional[...] for None-able types. See Item-class above, where the left and right neighbors can both be None.

Summary

So overall I was able to gain quite a speedup with an manageable amount of change to my existing code. In many places, I could improve the typing and clarity of code, so that was a nice side-effect too. Overall I loved using the familiar Python syntax and still being able to ramp up the runtime speed by an order of magnitude.

(By the way, Codon has nice Docs, so make sure to check them out!)

Thanks for reading and happy optimizing!