Python: type annotate class decorator return type

user_2476219 0
Viewed 106

How to annotate a class decorator that extends a class with additional members?

I have the following example where my class decorator inserts an async new function that calls the provided class' ainit. This code runs correctly, but the type returned by optimizer is the original type, and new is missing. How can I improve annotation, if possible?

AinitArguments = ParamSpec("AinitArguments")

class Optimizerable(Protocol):
    ainit: Callable[AinitArguments, None]


T = TypeVar("T", bound=Optimizerable)


def optimizer(cls: Type[T]) -> Type[T]:
    @classmethod
    async def new(cls: Type[T], *args, **kwargs) -> T:
        self = cls.__new__(cls)
        await self.ainit(*args, **kwargs)
        return self

    setattr(cls, "new", new)
    return cls

@optimizer
class Myclass:
    async def ainit(self, a: int, b: float) -> None:
        pass # stub
typing decorator python
1 Answers

Long story short, you can't really do this yet https://github.com/python/typing/issues/213

Python doesn't have a native type for an Intersection to represent a concrete class implementing two bases (either explicitly through inheritance or duck-typing like with protocol).

What you want would look like this (which just isnt possible yet):

class Optimizerable(Protocol):
    ainit: Callable[AinitArguments, None]

T = TypeVar("T", bound=Optimizerable)

class Optimized(Protocol[T]):
    async def new(cls: Type[T], *args, **kwargs) -> T:
        ...

def optimizer(cls: Type[T]) -> Type[Intersection[T, Optimized[T]]]:
    ...

I would either refactor this code to use a Mixin to add the method, or simply keep new as a stateless function, which avoids complex class structures entirely:

T = TypeVar("T", bound=Optimizerable)


async def optimized_new(cls: Type[T], *args, **kwargs) -> T:
    self = cls.__new__(cls)
    await self.ainit(*args, **kwargs)
    return self
user_32805380
Related