Decorator for forward type hint resolution

Question

Background

I frequently develop ORM libraries using the peewee framework. On one particularly large ORM library, representing a proprietary data format, I have a lot of coupled ORM models representing different nodes in an XML document. The library also provides de-/serialization methods for each ORM model from/to JSON and PyXB DOMs.

Since subordinate models need to known superordinate models for the foreign keys to be defined, I have this one-way dependency. Alas, also superordinate models need to know the subordinate models for the de-/serializations. To avoid putting all > 10,000 LOC into one humongous Python unit, I outsourced the DOM de-/serialization into mixin classes for each of the models into separate units, accessing the subordinate modules via the foreign keys of the deferred implementation, which avoids the cyclic imports.

I have not yet found a better solution to this problem. However, this means, that I do not have the implementated classes available on the respective module level, which again means that I cannot use them for type hinting. Hence, I use strings of the respective class names for type hinting:

# Class Foo not available here, since it implements FooMixin later.
class FooJSONMixin:
    @classmethod
    def from_json(...) -> 'Foo':  
        ...

However, this causes, that I dont't get the actual types in the method signatures in help() or pydoc, which I want.

Solution

Hence, I implemented a small library exposing a decorator to resolve type hint strings to the actual objects. I'd appreciate your feedback on this, since I deem it rather hacky.

"""Enhanced type hinting hacks."""

from typing import _UnionGenericAlias
from typing import Any
from typing import ForwardRef
from typing import Iterable


__all__ = ['resolve_type_hints']


def _resolve_union_generic_alias(
        uga: _UnionGenericAlias, mapping: dict[str, Any]
    ) -> Any:
    """Resolve the type."""

    return _UnionGenericAlias(
        uga.__origin__,
        tuple(_resolve_type_hint(a, mapping) for a in uga.__args__),
        inst=uga._inst,     # pylint: disable=W0212
        name=uga._name,     # pylint: disable=W0212
        _typevar_types=uga._typevar_types,  # pylint: disable=W0212
        _paramspec_tvars=uga._paramspec_tvars   # pylint: disable=W0212
    )


def _resolve_type_hint(value: Any, mapping: dict[str, Any]) -> Any:
    """Resolves a type hint."""

    if isinstance(value, str):
        return mapping.get(value, value)

    if isinstance(value, ForwardRef):
        return _resolve_type_hint(value.__forward_arg__, mapping)

    if type(value) is _UnionGenericAlias:   # pylint: disable=C0123
        return _resolve_union_generic_alias(value, mapping)

    return value


def _resolve_type_hints(
        annotations: dict[str, Any], mapping: dict[str, Any]
    ) -> None:
    """Resolves type hints of the given annotations dict."""

    for key, value in dict(annotations).items():
        annotations[key] = _resolve_type_hint(value, mapping)


def _resolve_mro(mro: Iterable[type], mapping: dict[str, Any]) -> None:
    """Resolve type hints of a method resolution order."""

    for cls in mro:
        for attribute in cls.__dict__:
            try:
                annotations = getattr(cls, attribute).__annotations__
            except AttributeError:
                continue

            _resolve_type_hints(annotations, mapping)


def resolve_type_hints(obj: type) -> type:
    """Decorator to resolves type hints on classes and functions."""

    mapping = globals()
    mapping[obj.__name__] = obj

    try:
        mro = obj.__mro__
    except AttributeError:
        _resolve_type_hints(obj.__annotations__, mapping)
    else:
        _resolve_mro(mro, mapping)

    return obj

Usage example

from typing import Optional, Union


class Model:
    pass


class TestMixin:
    @classmethod
    def foo(cls) -> 'Test':
        return cls()


@resolve_type_hints
class Test(Model, TestMixin):
    pass
    
    
@resolve_type_hints
def test(a: Optional['Test'], b: Union['Test', int]) -> 'Test':
    pass

Answer 1

I'm not sure I understood the problem you're trying to solve, but instead of writing:

class FooJSONMixin:
    @classmethod
    def from_json(...) -> 'Foo':
        ...

I wanted to point out that you can do:

from __future__ import annotations
from typing import TypeVar, Type


class BaseModel:
    ...


ModelT = TypeVar("ModelT", bound=BaseModel)


class ModelMixin:
    @classmethod
    def foo(cls: Type[ModelMixin]) -> ModelMixin:
        return cls()

At this point if you do something like:

class MyModel(ModelMixin, BaseModel):
    pass

Then MyModel.foo() return type will be MyModel.

Basically, you're telling mypy that whatever class the method foo is called from, it will always return an instance of this class.

Hope it can help