DEV Community

Cover image for Typing the untypable: generating Python .pyi stubs
Artem Vasilyev
Artem Vasilyev

Posted on • Originally published at Medium

Typing the untypable: generating Python .pyi stubs

In Python, some classes define their attributes dynamically. This means the attributes don't exist as normal instance variables and are often invisible to IDEs and static type checkers like mypy.

An abstract example:

class BookObject:
    fields = ("author", "title")

    def __init__(self, data: dict):
        self._data = {}
        for field in self.fields:
            self._data[field] = data.get(field)

    def __getattr__(self, name):
        if name in self._data:
            return self._data[name]
        raise AttributeError
Enter fullscreen mode Exit fullscreen mode

You can use it like this:

book = BookObject({"title": "Dune", "author": "Frank Herbert"})
print(book.title)
Enter fullscreen mode Exit fullscreen mode

It works at runtime, but your IDE will not autocomplete book.title, and mypy cannot check its type.

I saw this problem in oslo.versionedobjects  -  it is used in OpenStack for data models.

The library lets you define fields like this:

class Book(BaseObject):
    fields = {
        "id": fields.IntegerField(),
        "title": fields.StringField(),
        "author": fields.StringField(),
        "rating": fields.IntegerField(),
    }
Enter fullscreen mode Exit fullscreen mode

At runtime, you can access these fields like normal attributes (but there's a lot of magic happening under the hood), yet IDEs and mypy cannot see them.


Then I had an idea: for fun, I can add automatic .pyi stub generation to the project, so it is easier to work in IDE and mypy can find type errors.


.pyi files

Stub files (.pyi) are described in PEP 484.
They contain only type hints, and static analyzers use them.

Example:

class Book:
    title: str
    author: str
Enter fullscreen mode Exit fullscreen mode

Even if the real class attributes are dynamic, stub files let tools know what attributes exist and their types.

Demo project

I made a small demo project. Full code is here: demo repo.

Project structure:

.
├── main.py
├── Makefile
├── objects
│   ├── __init__.py
│   ├── base.py
│   └── book.py
├── pyproject.toml
└── stubgen.py
Enter fullscreen mode Exit fullscreen mode

Book class looks like this:

@DemoObjectRegistry.register
class Book(BaseVersionedObject):
    VERSION = "1.0"

    fields = {
        "id": fields.IntegerField(),
        "title": fields.StringField(),
        "author": fields.StringField(),
        "rating": fields.IntegerField(),
    }
Enter fullscreen mode Exit fullscreen mode

main.py has an example of its usage:

def echo(title: str, author_name: str, rating: str) -> None:
    print(f"{title.upper()} by {author_name.upper()} with rating {rating.strip()}!")

def main() -> None:
    books = [
        Book(id=1, title="The Doomed City", author="Arkady & Boris Strugatsky", rating=10),
        Book(id=2, title="The Settlement", author="Kir Bulychev", rating=10),
    ]

    for b in books:
        echo(b.title, b.author, b.rating)
Enter fullscreen mode Exit fullscreen mode

Notice: echo expects rating as str, but Book.rating is int.

Without stubs, mypy cannot detect this:

$ make lint
uv run ruff check .
All checks passed!
uv run mypy .
Success: no issues found in 5 source files
Nothing is detected because the dynamic fields are invisible.
Enter fullscreen mode Exit fullscreen mode

Generating .pyi stubs

To fix this, I wrote a small Python script stubgen.py that automates generating .pyi stubs with proper type hints.

Here is how it works.

Step 1 - generate baseline stubs

First, we generate initial stubs using mypy stubgen.

This produces .pyi files with classes and methods, but dynamic fields are missing:

class Book(BaseVersionedObject):
    VERSION: str
    fields: Incomplete
Enter fullscreen mode Exit fullscreen mode

Step 2 - parse source with ast and map field types

Then the script parses the original Python source with Python's ast module and looks for fields dictionaries in each class.

For each field, the script determines the Python type (int, str, etc.).

Each field constructor (like IntegerField or StringField) is mapped to a Python type:

FIELD_TYPE_MAP = {
    "IntegerField": "int",
    "StringField": "str",
    "BooleanField": "bool",
    "FloatField": "float",
}
Enter fullscreen mode Exit fullscreen mode

Step 3 - update the .pyi stubs

Finally, the script updates the .pyi stub by inserting typed attributes using ast:

# before
class Book(BaseVersionedObject):
    VERSION: str
    fields: Incomplete

# after
class Book(BaseVersionedObject):
    rating: int
    author: str
    title: str
    id: int
    VERSION: str
    fields: Incomplete
Enter fullscreen mode Exit fullscreen mode

Run stubgen and linters

Now we can generate the .pyi stubs by running:

$ make stubgen
Enter fullscreen mode Exit fullscreen mode

Next, we run the linters:

$ make lint
uv run ruff check .
All checks passed!
uv run mypy .
main.py:15: error: Argument 3 to "echo" has incompatible type "int"; expected "str"  [arg-type]
Found 1 error in 1 file (checked 5 source files)
make: *** [lint] Error 1
Enter fullscreen mode Exit fullscreen mode

Thanks to the generated stubs, mypy now detects that Book.rating is an int, while the function echo expects a string.

Conclusion

This is, of course, a simplified example - the script doesn't handle more complex cases like Optional, Dict, List, or nested objects.

But the main idea is that you can use this approach in your own projects to make working with dynamically defined attributes easier.

Top comments (0)