Enhancing CPython core with AST for Runtime Optimizations Using MyPy Pt. 1

I am earning no commission, whatsoever, by sharing any external links or referencing any link in this article and neither am I an affiliate or associated with any organization under whom falls the jurisdiction to govern these links.

source: https://unsplash.com/

Context

I recently went through few of USPyCon and EuroPyCon videos featuring Mark Shannon talking about his work on making CPython performant.

Other articles had mentioned this earlier:

Faster Python: Mark Shannon, author of newly endorsed plan, speaks to The Register

A Team at Microsoft is Helping Make Python Faster - Python

The presentation mentions profiling the slowest parts of the program using concepts like Adaptive Interpreter, Partial Evaluation, GC cycling, Memory management amongst others

If you have not seen one of these videos on this topic, it is highly recommended.

And then I went through this video (The interviewer, Michael, obviously feels overwhelmed here) where they discussed amongst other things, Mark's plans and other interesting ideas from the comments.

It was intriguing how many people have often discussed about integrating mypy or something similar more tightly such that type hints which are already sprinkled over python code can be generously utilized. To which, in the EuroPyCon talk at some point during the Q&A part of the presentation, Mark simply mentioned that statistical analysis, which Python is currently leveraging, is no worse than type hints so there does not seem to be significant work in this direction on his part.

Hm. I thought lets give some spare time to think in this direction. What has already been done before? What tools are availabel to leverage on this idea?

Research

I am no compiler theory expert but generally, it would help to understand how CPython compilation and execution flow works:

1. Parsing: The python script is parsed into an AST/Abstract Syntax Tree . Think of it like breaking down the syntax into its most fundamental tokens. The file you want to look under the hood here is `Python/ast.c` and header file `Include/internal/pycore_ast.h`
2. Compilation: This is where the parsed AST is broken down into bytecode. This is where these ASTs are translated into instruction sets that an interpreter can understand. Basically, the lingua franca devspeak code is translated to something that the interpreter can more efficiently translate to the machines lingua franca. To understand compilation, I would look at files that have to do with compilation i.e. the `Include/cpython/compile.h` header and `Python/compile.c`
3. Interpreting: This is where the bytecode is translated into execution paths. A Python Virtual Machine is the interpreter that translates intermetdiate representation of the code to the underlying hardware arhictecture which understands the OP codes. Here I googled around and found `Python/ceval.c` which executes the compiled code.
4. Note, that the compilation to bytecode can generates a .pyc file. This .pyc file is the compiled file which contains the bytecode that Python interpreter can execute. On subsequent imports of the same module, Python checks if a .pyc file exists and is up to date. If it is, it loads the bytecode from the .pyc file directly, which can speed up startup time since it skips the compilation step.

This makes sense. Since type hint checkers like ruff and mypy already do fair amount of type checks, is there a way we could leverage this to provide context to the compiler, since these tools have already done the first part i.e. parsing?

After a fair amount of googling into what existing tools have built upon these ideas, I did found out some solutions which I would like to highlight here but hang on to my flow of thoughts here. This is what I was thinking:

1. Introduce an optional static compilation step. If it can be made as part of the MyPy or Ruff tooling, the better. What I would ideally like is that no third-party tools need to involved here and devs could just do:

python compile [<script>.py]

And this would yield errors if any just like MyPy or ruff does. You can fix it and continue compiling or execute the code.

2. Ideally, the most straightforward way to do that is to generate the AST and pass it to compiler but with the metadata report of the types.

3. I had not looked under the hood to check what Python does regarding these types since types are inferred so I thought this would be the first thing to understand. How types are inferred in CPython?

4. The next obvious step would be to add this metadata information in the AST itself. By leveraging static type analysis and generating AST to optimize the bytecode compilation process in CPython as an optional compilation step without altering the dynamic nature of the language, it would require considerable tweaking in the implementation files written in C.

Which means it will mess up the entire chain of command all the way from compilation to Python VM. If PSF steering committe ever thought of that, they would have realized this is probably the last ditch effort to boost the runtime performance. That probably is what Mark said what he said in the videos aforementioned on why it has not made sense to make an effort here.

But If there anymore juice to squeeze, this is probably one area that looks reasonable from the outside i.e. as python developers this makes sense. Since I was enjoying tinkering lately with ASTs too much in other languages and they have wrapped my head in a rather fractal-convoluted way (If you have one of those mental conditions, do not click).

What Work Exists Already?

So back to the tools that already exist.

Let's go back to what are the requirements here by skimming over this snippet:

source_code = """
def greet(name: str) -> str:
    return "Hello, " + name
"""

parsed_ast = ast.parse(source_code)
enhanced_ast = enhance_ast(parsed_ast, source_code)

What I would like to have is to do this:

python compile [<script>.py]

which generates some output that can then be optionally fed into the actual compilation procedure and simply run:

python <script>.py

Aha! MyPy has stubgen which does this exactly. Initially I tried this:

mypy script.py > mypy_output.txt 2> mypy_errors.txt

But unless you have not fixed the issues yourself, which is what mypy is there to warn you in the first place, this would produce a no issues found or a success message in the output and no errors, which should be our starting point.

What you really want is this:

stubgen hello_world.py

for a script to testlike:

def hello_world() -> None: 
    print("Hello, world!")

def add2num(a: dict[str, int], b: dict[str, int]) -> int:
    return a.get("a", 0) + b.get("b", 0)

that gives you a pyi file like:

def hello_world() -> None: ...
def add2num(a: dict[str, int], b: dict[str, int]) -> int: ...

There have been several significant efforts within the Python community aimed at integrating static type analysis with AST manipulations, which align with the concepts discussed here such as:

1. Pytype by Google: This tool performs type checks without requiring type annotations. It can infer types even when they are not specified, analyze code for type errors and generate type annotations. This could be useful in enhancing ASTs with inferred type information, although it focuses on linting and error detection rather than directly optimizing AST for runtime performance. However, I did not want to jump the ship yet with inference. I first want to inject the already considered and generated type. So I will bypass this for now.
2. Scalpel: Scalpel is a static analysis framework designed for Python that provides capabilities such as rewriting code for better static analysis and optimizing problematic programs. This is about enhancing the code and its signature from what I understood. Their documentation states Detecting bugs, Fixing vulnerabilities, Profiling code and Refactoring code
3. MyPy: Then there is mypy itself, which provided the insight I was looking for to begin tinkering something.

Once you have this:

def hello_world() -> None: ...
def add2num(a: dict[str, int], b: dict[str, int]) -> int: ...

What I am looking for is this:

def enhance_ast_with_stubs(source_code_path, stubs_path):
    # Read the original script and generated stub
    with open(source_code_path, "r") as source_file, open(stubs_path, "r") as stubs_file:
        original_ast = ast.parse(source_file.read())
        stub_ast = ast.parse(stubs_file.read())
        # Enhance the ast with type information
        enhanced_ast = attach_type_annotations(original_ast, stub_ast)
        # ast_dump = ast.dump(enhanced_ast)
        # print(ast_dump)
        print("The ast unparsed")
        print(ast.unparse(enhanced_ast))
        return enhanced_ast

In the followup article, I will conclude this idea with my exploration further which would specifically require me forking the latest CPython in my spare time and tinkering on my local to mess things with. Chime in as I would like to understand what insights can you provide here. There is always reddit and stackoverflow to Q&A things around this exploratory phase.

👋 If you find this helpful, please click the clap 👏 button below a few times to show your support for the author 👇

🚀Join FAUN Developer Community & Get Similar Stories in your Inbox Each Week