CODEWITHSUNDEEP
pythondecoratortyping
Aart Stuurman
Aart Stuurman

Reputation: 3628

Python: type annotate class decorator return type

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

Upvotes: 2

Views: 408

Answers (1)

flakes
flakes

Reputation: 23644

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

Upvotes: 1

Related Questions