This is a continuation of my previous post on this topic. I’ve decided to make this a new thread since the ideas have been refocused a lot since then and my goal now is to essentially make a new run down of the problem and my proposed solution since there was a lot of confusion before. This post won’t contain every detail of my ideas and implementation, I’ll be more comprehensive in the rewrite of my draft PEP, but I hope it’s understandable.
Motivation
Just to recap, the way that type annotations currently work is that when the compiler encounters a function definition like def func(a: int, b: some_module.MyClass): ... it not only produces the bytecode for func itself, but also for an annotate function that is then stored as func.__annotate__ and behaves essentially like this:
def __annotate__(format, /):
if format > Format.VALUE_WITH_FAKE_GLOBALS:
raise NotImplementedError
return {
"a": int,
"b": some_module.MyClass,
}
This function is then later called to e.g. create the __annotations__ dict or by tools like the dataclass decorator. What this implementation solves are issues like forward references (since the function can be called later when the references are resolved) or not being able to properly inspect annotations at runtime (since the returned dict will contain the exact objects the annotation expressions specify and things like name lookups are resolved properly through the usual machinery).
But this implementation also means that annotation expressions are evaluated using the exact same semantics that normal expressions occurring in other places are. This heavily limits what kinds of expressions we can sensibly use as type expressions. For example, one might reasonably expect that Python uses syntax similar to typescript when defining literal types, you just write down the values like this: error_level: 1 | 2 | 3. But that’s not possible since at runtime that annotation looks identical to writing error_level: 3 because | is treated as the “bitwise or” operation and everything is constant folded together. So we have to write error_level: Literal[1, 2, 3] to make sure each literal value is preserved.
This restriction has now come up in several discussions. This article discusses several common wishes users have for writing types in an easier way, most of which are (in part) blocked by this. There has been work on type checking array sizes, one of the conclusions of which was that the syntax that is currently possible is too cumbersome for them to see real-world use. A long-standing wish is to have inline typed dicts, which currently would be forced to use needlessly cumbersome syntax. And most recently, PEP 827 proposes a comprehensive type manipulation system, which (IIRC) originally proposed a set of new type forms to handle this issue, but then expanded to a system similar to my proposal since it was deemed too unreadable.
I want to be clear that I am not proposing to implement any one of these examples or that the only thing making them impossible or difficult is how annotations are evaluated. Rather, my goal is to implement the baseline feature that makes additions like these possible. I think that while any one of these proposals might not meet the bar to be implemented, the fact that we now have several ideas that are promising and cause a lot of discussion that all need some variant of this change is a strong argument for its adoption.
My Proposal
My idea is to change annotate functions to no longer contain the bytecode that produces the value that expressions are evaluated to, but rather some data objects that encode the original expression AST. The above example could then look something like this (when only considering the code path for the AST format):
def __annotate__(format, /):
asts = {
"a": _build_ast_object("...some data..."),
"b": _build_ast_object("...some data..."),
}
namespace = {}
try:
namespace["int"] = int
except NameError:
pass
try:
namespace["some_module"] = some_module
except NameError:
pass
return asts, namespace
The AST data itself is stored in some opaque binary blob and is transformed into the normal ast.AST objects at runtime. We also compute a namespace that contains every name referenced in the annotation expressions. These are later used to evaluate the actual type objects that the annotations are defining and the try/except handling prevents forward references from being an issue.
This evaluation is done by some new helper functions added to the typing module. These behave similar to annotationlib.get_annotations. But they don’t directly pass through the return value of the annotate functions and instead call them to get the AST objects and then evaluate these using the returned namespaces and typing-specific semantics.
This implementation preserves backwards-compatibility since calling annotate functions with any of the existing formats returns the exact same objects. The possible new semantics that are specific to typing uses are only produced when you use the new typing helpers.
One downside of this solution is that this machinery is limited to annotate functions, and it thus doesn’t allow you to use any newly implemented typing features in some places. In particular, things like cast(Some Fancy Type, ...) will still be evaluated using the usual semantics and thus produce errors or even incorrect values. However, we can easily work around this by using type aliases like this:
type _MyAlias = Some Fancy Type
...
cast(_MyAlias, ...)
It’s certainly not a perfect solution, but since annotate functions cover most places where users actually want to express types I don’t think it’s too much of a hassle to sometimes have to fall back to this. And since type checkers can detect places where issues might occurr and the vast majority of users that are writing type expressions also run type checkers, accidentally forgetting an intermediate type alias shouldn’t happen too often.
I’ve implemented this functionality and several alternatives in my CPython fork. Other variants are implemented in other branches, but a lot of the approaches that didn’t pan out well are lost somewhere in the git tree. I can try to collect a full set of variations if wanted.
Alternative Implementations
Another proposed approach is to not produce AST objects, but rather to just store the annotation objects as strings, similar to PEP 563. I’ve tried out various implementations using both ASTs and strings and did some (admittedly basic) performance measurements. I measured memory usage by importing a selection of the most popular packages on PyPi that use type annotations and observing total used memory. The vast majority of that obviously are not the annotate functions themselves, but the differences between implementations were consistent enough that we can use these measurements comparatively. I then also measured the total execution time of the annotate functions in these modules. These are my relative results for each of the best variants:
| Current | AST | Strings | |
|---|---|---|---|
| Memory | 100% | 99.4% | 99.2% |
| Time | 1x | 7x | 24x |
While the memory numbers are great, the runtime measurements are much, much less exciting. There unfortunately just isn’t really a way (or at least I haven’t found it) to get around the fact that executing bytecode is pretty fast compared to analzying some complex structure. To put some further context to these numbers, the difference in RAM usage between the AST and strings based approaches is within the range of my codegolfing. On this front, all three approaches are more or less identical. But the runtime measurements had much larger differences than simple omptimizations. Neither approach is great, but storing the AST data directly has the huge advantage that you then don’t need to parse the strings at runtime.
Because of these results, I heavily favor the AST based approach. It does have some downsides such as implementation complexity, but it performs much better and is more expandable in the future. In particular, I’ve found that skipping the resource heavy ast.AST objects as they are currently implemented and/or constructing the type objects in C cuts down execution time by about half. While either of those optimizations have significant disadvantages, it does mean that we have a lot more room for future improvements in the internal processing. Storing the annotations as strings locks us into a much more rigid representation and pipeline.
I also don’t think that these numbers should discourage us from this proposal. Slowing execution down by a factor of 7 is bad, but (at least in my observations) code generally doesn’t spend that much time constructing annotation objects for that to be a significant performance loss. And should that estimation be incorrect, we can save a lot of that overhead using different implementations that follow the same external API.
Open Questions
There are various open questions and implementation choices that I’m not sure about:
As was recently mentioned, associating a unique namespace to every annotation rather than returning a single namespace from the annotate function makes merging annotations from different modules more straightforward. It also keeps the annotation protocol more consistent since annotate functions then always return a dict mapping parameter/attribute names to values (or a single such value) instead of the AST format returning a tuple of such a dict and a namespace. However, this also duplicates a lot of information.
How do we want to handle forward references? My current implementation just builds a dict of existing names, which means that forward references are simply missing from the namespace. If you want to later resolve that forward reference, you have to call the annotate function again. This could be changed by not using a namespace dict but a function closure, but that increases code size.
How exactly do we store the AST data? My currently favored approach stores a binary representation of a DFS traversal of the AST nodes in a str object. This lets us greatly benefit from string interning, but also means that we need to maintain a C code serializer and parser for this. The vast majority of that code can be automatically generated from the grammar, but it’s still a somewhat finnicky format.
There’s probably many more things that still need to be figured out and things I missed here. I’m happy to get any feedback and thoughts.
