Building Robust Codebases with Python's Type Annotations
Hudson River Trading's (HRT's) Python codebase is large and constantly evolving. Millions of lines of Python reflect the work of hundreds of developers over the last decade. We trade in over 200 markets worldwide — including nearly all of the world's electronic markets — so we need to regularly update our code to handle changing rules and regulations.
Our codebase provides command-line interface (CLI) tools, graphical user interfaces (GUIs), and event-triggered processes that assist our traders, engineers, and operations personnel. This outer layer of our codebase is supported by an inner layer of shared business logic. Business logic is often more complicated than it appears: even a simple question like "what is the next business day for NASDAQ?" involves querying a database of market calendars (a database that requires regular maintenance). So, by centralizing this business logic into a single source of truth, we ensure that all the different systems in our codebase behave coherently.
Even a small change to shared business logic can affect many systems, and we need to check that these systems won't have issues with our change. It's inefficient and error-prone for a human to manually verify that nothing is broken. Python's type annotations have significantly improved how quickly we can update and verify changes to shared business logic.
Type annotations allow you to describe the type of data handled by your code. "Type checkers" are tools that reconcile your descriptions against how the code is actually being used. When we update shared business logic, we update the type annotations and use a type checker to identify any downstream systems that are affected.
We also thoroughly document and test our codebase. But written documentation is not automatically synchronized with the underlying code, so maintaining documentation requires a high level of vigilance and is subject to human error. Additionally, automated testing is limited to the scenarios that we test for, which means that novel uses of our shared business logic will be unverified until we add new tests.
Let's look at an example of type annotations to see how they can be used to describe the shape of data. Here's some type annotated Python code that computes the checksum digit of a CUSIP:
def cusip_checksum(cusip8: str) -> int:
assert len(cusip8) == 8
chars: str = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ*@#"
charmap: dict[str, int] = {
char: value
for value, char in enumerate(chars, start=0)
}
total: int = 0
for idx, char in enumerate(cusip8, start=0):
value: int = charmap[char]
if (idx % 2) == 1:
value *= 2
total += (value // 10) + (value % 10)
return (10 - total % 10) % 10
Here's what the type annotations tell us:
cusip_checksum()is a function that takes a string as input and returns an integer as output.charsis a string.charmapis a dictionary with string keys and integer values.totalandvalueare integers.
HRT uses mypy to analyze our Python type annotations. Mypy works by analyzing the type annotations in one or more Python files and determining if there are any issues or inconsistencies.
Most of the time, mypy is good at type inference, so it’s better to focus on annotating the parameters and return values of a function rather than the internal variables used in a function.
Here's a new function, validate_cusip(), that relies on the cusip_checksum() function from earlier:
def cusip_checksum(cusip8: str) -> int:
...
def validate_cusip(cusip: str) -> str | None:
checksum: int
if len(cusip) == 9:
checksum = cusip_checksum(cusip[:8])
if str(checksum) == cusip[8]:
return cusip
else:
return None
elif len(cusip) == 8:
checksum = cusip_checksum(cusip)
return f"{cusip}{checksum}"
else:
return None
Mypy is happy with this code:
Success: no issues found in 1 source file
Now, let's say that we decide we should update cusip_checksum() to return None if it detects that the CUSIP is not valid:
def cusip_checksum(cusip8: str) -> int | None:
if len(cusip8) != 8:
return None
chars: str = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ*@#"
charmap: dict[str, int] = {
char: value
for value, char in enumerate(chars, start=0)
}
total: int = 0
for idx, char in enumerate(cusip8, start=0):
try:
value: int = charmap[char]
except KeyError:
return None
if (idx % 2) == 1:
value *= 2
total += (value // 10) + (value % 10)
return (10 - total % 10) % 10
Mypy automatically detects issues in how validate_cusip() is using cusip_checksum():
error: Incompatible types in assignment (expression has type "int | None", variable has type "int") [assignment]
Now that we've been alerted, we can update validate_cusip() to handle these changes:
def cusip_checksum(cusip8: str) -> int | None:
...
def validate_cusip(cusip: str) -> str | None:
if len(cusip) == 9:
match cusip_checksum(cusip[:8]):
case int(checksum) if str(checksum) == cusip[8]:
return cusip
elif len(cusip) == 8:
match cusip_checksum(cusip):
case int(checksum):
return f"{cusip}{checksum}"
return None
In this example, the functions were next to each other in the source code. But Mypy really shines when the functions are spread across many files in the codebase.
All in all, type annotations have substantial benefits for making your codebase more robust. They are not an all-or-nothing proposition — you can focus on adding type annotations to small parts of your codebase and growing the amount of type annotated code over time. Along with other technologies, Python's type annotations help HRT to continue thriving in the fast-paced world of global trading.
This article originally appeared in the HRT Beat.
Meet the author:
John Lekberg works on a spectrum of Python and gRPC systems at HRT. He primarily develops and refines internal tooling for monitoring and alerting. He's also led initiatives to apply static analysis tools to HRT's codebase, catching bugs and reducing the manual work needed to review code.