[Python] Python type-annotation for getitem

I'm implementing a custom class and want to provide type hints for item access (e.g. obj[key]) so that

• Pyright/Pylance/VSCode suggests allowed item keys,
• Pyright/Pylance/VSCode suggests correct methods on the return values, and
• MyPy performs type checks when the returned values are used.

When using the custom class, the user should provide a schema that indicates allowed item keys and their types, analogously to TypedDict.

In this post I describe the problem very generically. My specific motivation is to extend type annotations for DataFrames so that column access via __getitem__ is aware of the columns type as this would greatly improve maintainability of data science code heaviliy relying on DataFrames. I'm aware of Pandera. It does provide DataFrame schemas, but doesn't really support column access. I'm also aware that the exact columns can be different at runtime, but that applies to all type hints in Python.

Attempt 1: Extending TypedDict


TypeDict provides the features I'd like to have with my own class. So my first idea was to inherit from TypedDict. However, Python doesn't allow inheriting from TypedDict and other classes. Not being able to inherit from other classes is a too big limitation.

from typing import TypedDict

class MyClass(TypedDict, OtherClass):
x: str
y: int

Attempt 2: Providing Stubs


I managed to get correct coding suggestions in VSCode by providing stubs.

# myclass.pyi

from typing import Literal, overload


class MyClass:
@overload
def __getitem__(self, key: Literal["x"]) -> str: ...

@overload
def __getitem__(self, key: Literal["y"]) -> int: ...


# myclass.py

class MyClass:
def __getitem__(self, key):
if key == "x":
return "example string"
elif key == "y":
return 42



With this, I get suggestions for the allowed keys and their items.

var = MyClass()

# Dict keys and methods suggested by Pylance
print(var["x"].capitalize())
print(var["y"].as_integer_ratio())


However, using them incorrectly isn't spotted by any of the tools.

# No problem reported by Pylance or mypy
print(var["y"].capitalize())
print(var["x"].as_integer_ratio())


It is very cumbersome to define stubs. In the context of DataFrames, users would need to define stubs every time they define a schema. That is a lot of boilerplate code. It would be great if the stubs could be created dynamically.

Attempt 3: Meta programming stubs


The idea is to use a meta class that creates the stubs automatically. However, I don't know how to create the overloaded methods as they cannot coexist in the objects __dict__. I'm also not sure if my meta class would even run during type checks.

Question

• How to get mypy to report type errors with my stubs?
• Is there a way to dynamically create stubs based on a schema?
• Are there better alternatives?
• How does TypedDict work in type checks? I TypedDict implemented using standard type hints or is its behavior hard-coded?

Continue reading...