How can I use generic classes with a generic function in Python?

user_15698819 0
Viewed 429

I had trouble figuring out how to phrase my question, but I have code like the following:

  1. There's a generic class (BaseOuter) with a method (to_inner) whose return type is generic.
  2. There are concrete subclasses of that generic class (e.g. Outer). When I call Outer().to_inner(), the inferred return type is Inner as I expect.
  3. I create a generic function that takes an optional BaseOuter and calls to_inner() on it. But when I pass an Outer to this function, the inferred return type is BaseInner, and not Inner as I expected.

Am I doing something wrong? Is there a way to make type checkers realize that to_inner(Outer()) returns an Inner?

import typing

class BaseInner:
    pass

T_co = typing.TypeVar('T_co', bound=BaseInner, covariant=True)
T = typing.TypeVar('T', bound=BaseInner)


class BaseOuter(typing.Generic[T_co]):
    def to_inner(self) -> T_co:
        return self.InnerType()


@typing.overload
def to_inner(outer: BaseOuter[T]) -> T: ...
@typing.overload
def to_inner(outer: None) -> None: ...
def to_inner(outer: typing.Optional[BaseOuter[T]]) -> typing.Optional[T]:
    return outer.to_inner() if outer else None


class Inner(BaseInner):
    pass


class Outer(BaseOuter[Inner]):
    InnerType = Inner


this_is_an_Inner = Outer().to_inner()
unfortunately_a_BaseInner = to_inner(Outer())
python-typing python
1 Answers

You need to add InnerType annotation to BaseOuter generic class, because this name is used in to_inner method. Otherwise, mypy (0.910) will not be able to understand your intentions:

class BaseOuter(typing.Generic[T_co]):
    InnerType: typing.Type[T_co]   # <---
    
    def to_inner(self) -> T_co:
        return self.InnerType()
user_137826690
Related