Writing a Python Type Checker Pt. 1

Lately I have been seeing a bit of excitement online around Python type checkers. One of my biggest frustrations with Python is dealing with unknown types and having them blow up at run time. I have been aware of these tools, and used mypy a little to help with this.

With Python being the mess of types that it can be I have been curious how type checkers for Python actually work. It seems like an overwhelming task to take on. After hearing about new and in progress tools like pyrefly from Meta and ty from Astral, I figured these projects might be approachable enough right now to dig into the code base and learn about how they are implemented.

This series I’m writing goes through writing a type checker for Python based on what I’ve figured out. It will be really simple. The goal for me is to get an idea of how they work and get more familiar with type checking in general, not develop a tool that can compete with some of these big projects with teams of experts.

The project is in progress as I write this. The code is at: https://github.com/tejom/python-type-check

Getting Started

I needed to pick a language to write this in so this series will be implementing the checker in Rust. Pyrefly and Ty are both also implemented in Rust, so it helps to read the code and compare. I also think Rust’s features make some of the language related work pretty neat.

I’ll mostly focus on type checking functions, integers and strings. The program will infer types and report errors. I liked how Ty reported errors with context so I’ll reuse the format by the end.

Example from the python-type-checker:

[Error] src.py:7:4 Invalid types (Integer(1), String) for BinOp
 5 | a = "foo"
 6 | f = x + y
 7 | b = 1 + a
   |     ^^^^^

I want this project to be about type checking and not everything language related. For the AST, instead of writing a parser, I am using tree-sitter and tree-sitter-python for the grammar.

Walking the AST

At a real simple level type checking seems to be, track a variable’s type as a program is run. To run a program we need to walk the AST. I wanted to set up a way to visit a node in the tree and do something with it.

The first thing I wrote is a pretty printer for the AST. This does two things for me, sets up a generic way to take actions on each Node of the tree and helps visualize the nodes and their children as each node is visited.

Tree sitter uses a cursor, a stateful struct that points to the current node, to move through the tree. To do something like DFS, I’m using recursion and calling a visit children function. The language’s stack will act as the stack instead of pushing Nodes onto an actual stack data structure. The programs are mostly simple enough that it shouldn’t hit a stack overflow with this approach.

The recursive function’s signature for walking the tree is

visit_all_children(cursor: &mut TreeCursor, visit_cb: &mut dyn FnMut(&mut TreeCursor))

When each node is visited, a call back will be called, then it will go through all of the children and call the function again with the call back.

“Pretty Printer” will be simple to write with this recursive function. For each node, call the print node call back function, which prints the node info indented by the depth of the node in the tree. That info is contained in the Node struct.

For the python program:

a = 5
b = 3
c = a + b

The AST that is printed:
Each line is [“node type” (node text) depth]

  "expression_statement" (a = 5) 1
   "assignment" (a = 5) 2
    "identifier" (a) 3
    "=" (=) 3
    "integer" (5) 3
  "expression_statement" (b = 3) 1
   "assignment" (b = 3) 2
    "identifier" (b) 3
    "=" (=) 3
    "integer" (3) 3
  "expression_statement" (c = a + b) 1
   "assignment" (c = a + b) 2
    "identifier" (c) 3
    "=" (=) 3
    "binary_operator" (a + b) 3
     "identifier" (a) 4
     "+" (+) 4
     "identifier" (b) 4

Starting the type checker

The type checker will also visit every node with a callback. Instead of printing the node, it will type check the node and track type information.

Check_visit Will be the type checker’s call back function. Each node type will be handled differently.

First type errors to detect

The first errors I’ll look into detecting are passing mismatched types into “binary_operator” nodes

The python script was something like:

a = 1
b = "abc"
x = a + b

Tracking Types

Pyrefly’s README actually had a nice little example of how to approach this. I also read through ty’s Usedef comments and a few other files in the repo.

The approach I ended up with in the beginning was tracking locations(Place struct) with identifiers in the file and mapping those to an inferred type. This I called “bindings”, implemented with a hash map. Then with a second hash map, track the current live variable to a point in the source code. A variable’s type can be looked up by finding the place in the source code that currently represents it and then finding that location’s type

Place("a @ line 2, col 2") => Type::Int
and
Identifier("a") => Place("a @ line 2, col 2")

I eventually refactored these hash maps into an Environment struct that held a stack of “Scopes” when I moved on to handling functions

Tracking Assignment

Before getting to operations like addition, I needed to store the variable types used.

A = 1 is an expression, it is assignment and has two child nodes in the AST, one for left, one for right.

Assignment tracks the variable on the left by storing the place it’s located and adding a value for the variable name in the hash map.

The left side’s type is determined by how ever the right side ends up being inferred. There’s a function infer type for expression that then goes into the right side and figures out what type it evaluates to. This type and the left side’s place is then stored in the bindings hash map. Assignment makes that value the live value so the hash map of variables to locations is updated too.

Checking Binary Operations

I was tricked into thinking `+` would end up being simple, because addition is simple.

The goal is check if this code has the correct types:

1 + 4
a + b

And determine if the output can be assigned correctly

z = 1 + 5
reveal_type(z) # should be int

What ended up being tricky is that there is some kind of implied signature for addition and binary operations in general. I thought something like `T + T -> T` maybe. I got worried that I would need to implement some kind of generic type check and a way to infer and handle them. After taking a look at how Ty handled this, I saw they have a hard coded list of types for signatures and return values and figured this would be good enough for me for now.

Reading through this I also saw that addition was a bit more complicated because of how the function is found. `+` is syntactic sugar for a class’s __add__ function. So to type check addition in general you need to look at the class, then get the method. __add__ for numbers doesn’t seem to have a typed signature it would need to be special cased anyway.

If I added more operations like `*` then `str * str -> str` would be valid, but so is `str * int -> str`

Check_binop infers the types for the left and the right then determines the return value. The return value can be used to infer the type if it is used in assignment.

Whats next

Part 2 will move on to checking function definitions, calling functions and implementing reveal_type

I might go on to handle If’s, generics, and protocols if I have the time

The repo is for learning, if you stumble upon this and want to try hacking an implementation for something or practice working with Rust feel free and maybe even open a PR!