Little Python trick: Type-safe object extensions

I’m currently playing with type annotation and static type checking in Python (using mypy) and came across a somewhat interesting problem: How to allow extending an object in a type-safe manner?

Suppose that you have an object where extension modules can add attributes: for instance, the Django Web framework’s sessions middleware adds a session attribute on each incoming request object so that views can use the session.

This approach does not work with static type checking: For the the type checker, a Request object has a fixed number of attributes, and session is not one of them. The session attribute cannot be hard-coded into the request object because anybody can write a new middleware that will also have this problem.

A solution can be to add a generic dictionary to the Request object so that extensions can add their data. The values of this dictionary cannot be bound to a single type, since each extension will need to add its own type of data, so it will be typed Any:

class Request:
	ext_values: dict[str, Any]

Now the session middleware can do request.ext_values['session'] = … and the view code can use it without issues. However, for the type checker, request.ext_values['session'] is an untyped (Any) object, which means that it cannot do any typing verification on the object. Typos like request.ext_values['session'].usre will not be detected.

The view code can also cast the session object from Any to its actual type, but that’s not very pretty, and could be error-prone.

So how can we get the type checker to understand “If I access this dictionary with this key, it will always return an object of this exact type”? By using the type as a dictionary key!

T = TypeVar('T')

class Request:
    _ext_values: dict[type, Any] = {}

    def get_ext(self, val_type: Type[T]) -> T:
        return cast(T, self._ext_values[val_type])

    def set_ext(self, value: object) -> None:
        self._ext_values[type(value)] = value

This may require a little explanation if you are not familiar with Python’s metaprogramming capabilities. If you write class A:, the expression A is itself an object (usually an instance of the type class), and this object work like most objects. Most notably for our case, we can compare classes (they are all different from each other) and they have a hash value, so they can be used as dictionary keys:

class A:
 pass

class B:
 pass

d = {
 A: 'a'
}

d[B] = 'b' # d now contains 2 values
print(d[A], d[B]) # prints 'a', 'b'

So now let’s see what happen on the Request object:

# Session middleware code
session = Session(…)
request.set_ext(session)

set_ext calls type() on the provided session object, which returns the class object (here, Session), and uses it as a key to store the session object. Now we have _ext_values == {Session: session}.

In the view code:

user = request.get_ext(Session).user

Here, the Session class is passed to get_ext. Since it exists in the dictionary, get_ext returns the stored session object, and user code can retrieve the user object inside it. The important parts are the type annotations: the -> T means “I return an object of generic type T” and the : Type[T] means “I take an argument which is a class object, and this class object is T”

This means that when the type checker sees the request.get_ext(Session) code, it know that this expression returns an instance of Session, and can check that it actually has a user attribute. Type safety achieved!

One cool thing is that despite the cast done on object coming out of the dictionary, this is type safe: as long as nobody messes up with _ext_values directly, get_ext will always return an instance of the type passed in, or a KeyError. Other extensions cannot accidentally erase a value set by another extension (as long as they use distinct types).

Note that it is pretty simple to extend this scheme to multiple values per class, still in a type-safe manner: add a string key, and use a (type, string key) tuple as the dictionary key:

class HTTPRequest:
    _ext_values: dict[tuple[type, str], Any] = {}

    def get_ext(self, val_type: Type[T], key: str) -> T:
        return cast(T, self._ext_values[(val_type, key)])

    def set_ext(self, key: str, value: object) -> None:
        self._ext_values[(type(value), key)] = value

Now the session code does request.set_ext('session', Session(...)) and the view code does request.get_ext(Session, 'session'). The use of both the type and the string key to index values in the dictionary means that type safety is preserved: If another extension does `request.set_ext('session', AnotherObject(...)), the dictionary will contain:

{
	(Session, 'session'): <the session object>,
	(AnotherObject, 'session'): <the other object>,
}

and request.get_ext(Session, 'session') will still return the session object.

Note that I didn’t invent this technique; it’s rarely used in the wild, because it requires a programming language where types are first-class objects (so we can use them as dictionary keys) and that perform static type checking. I’m pretty sure I’ve seen it used in Swift (which has both characteristics).