Python typing hell - decorated descriptors

type_rewriter.py

from __future__ import annotations

from collections.abc import Callable
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Concatenate, Generic, ParamSpec, TypeVar, cast, overload

_T = TypeVar("_T")
_F = TypeVar("_F", bound=Callable[..., Any])
_P = ParamSpec("_P")
_R_co = TypeVar("_R_co", covariant=True)


class AnnotatedMethod(Generic[_T, _P, _R_co]):
    _f: Callable[Concatenate[_T, _P], _R_co]
    _mod: str
    _qn: str

    # FIXME: Need weakref?

    def __init__(
        self, func: Callable[Concatenate[_T, _P], _R_co], module: str, qualname: str
    ) -> None:
        self._f = func
        self._mod = module
        self._qn = qualname

    @overload
    def __get__(self, obj: None, cls: type[_T], /) -> Callable[Concatenate[_T, _P], _R_co]: ...
    @overload
    def __get__(self, obj: _T, cls: type[_T] | None = None, /) -> Callable[_P, _R_co]: ...
    def __get__(
        self, obj: _T | None, cls: type[_T] | None = None, /
    ) -> Callable[Concatenate[_T, _P], _R_co] | Callable[_P, _R_co]:
        if obj is None:
            return self._f
        return cast(Callable[_P, _R_co], self._f.__get__(obj, cls))

    def __set_name__(self, obj: Any, name: str) -> None:
        if obj is None:
            raise ValueError(f"None obj? {obj}")
        if not hasattr(obj, "_infos"):
            setattr(obj, "_infos", {})
        obj._infos[self._mod, self._qn] = {"self": self, "name": name}


class rewriter:
    _mod: str
    _qn: str

    def __init__(self, module: str, qualname: str) -> None:
        self._mod = module
        self._qn = qualname

    def __call__(self, func: _F) -> _F:
        return cast(_F, AnnotatedMethod(func, self._mod, self._qn))


class TypeRewriter:
    _infos: dict[tuple[str, str], Any]
    _infos_ro: MappingProxyType[tuple[str, str], Any]

    def __init__(self) -> None:
        self._infos_ro = MappingProxyType(self._infos)

    @rewriter("typing", "Union")
    def fancy(self, a: int, b: int) -> int:
        print(f"fancy() self: {self} a: {a} b: {b} infos: {self.infos}")
        return a + b

    @rewriter("pycparser.c_ast", "Union")
    def mancy(self, a: int, b: int) -> int:
        print(f"mancy() self: {self} a: {a} b: {b} infos: {self.infos}")
        return a * b

    @property
    def infos(self) -> MappingProxyType[tuple[str, str], Any]:
        return self._infos_ro


if __name__ == "__main__":
    tr = TypeRewriter()
    print(f"tr.fancy(1, 2): {tr.fancy(1, 2)}")
    print(f"TypeRewriter.mancy(b, 7, 11): {TypeRewriter.mancy(tr, 7, 11)}")