Possible to specify generic typed class via its __init__()?

[abelcheung]

I'd want to write a stub for some external class (not controlled by me), where some of its method may return different type depending on how it is constructed. For example:

_T = TypeVar('_T', str, int)

class Klass(Generic[_T]):
    def __init__(self, is_int: bool) -> None: ...
    def common_method(self) -> bool: ...
    def do_it(self) -> list[_T]

x = Klass(True)  # x expected to be Klass[int]
y = Klass(False)  # y expected to be Klass[str]
x.do_it()  # expects list[int]
y.do_it()  # expects list[str]

Is such maneuvering possible without forcefully annotating x and y? (like x: Klass[int] = ...)

[BvB93]

An option here would be to overload __init__ w.r.t. boolean literals:

class Klass(Generic[_T]):
    @overload
    def __init__(self: Klass[int], is_int: Literal[True]) -> None: ...
    @overload
    def __init__(self: Klass[str], is_int: Literal[False]) -> None: ...

[abelcheung]

Thanks a great deal, I have never thought of annotating self as well.

[erictraut]

A word of caution here... Annotating the self parameter of an __init__ method to indicate the type of the constructed class is not documented or specified anywhere. It's something that mypy introduced and other type checkers have implemented to varying degrees. There are a bunch of edge cases with this technique that do not work. You should be OK if the type annotation involves a concrete specialization of the containing class as in the example above. If you use type variables, unions, or other classes, your results will vary.

An alternative solution that doesn't rely on this unspecified behavior is to provide a __new__ method rather than an __init__.

class Klass(Generic[_T]):
    @overload
    def __new__(cls, is_int: Literal[True]) -> Klass[int]: ...
    @overload
    def __new__(cls, is_int: Literal[False]) -> Klass[str]: ...