Making a subclass of a non-generic subclass of ContextManager more specific

[OddBloke]

Given the following classes:

from typing import ContextManager

class SuperClass(ContextManager[SuperClass]): ...
class FirstSubClass(SuperClass): ...
class SecondSubClass(SuperClass, ContextManager[SecondSubClass]): ...

a reveal_type for each C.__enter__ gives:

cm.py:24: error: Revealed type is 'def () -> cm.SuperClass*'
cm.py:25: error: Revealed type is 'def () -> cm.SuperClass*'
cm.py:26: error: Revealed type is 'def () -> cm.SuperClass*'

What I'm trying to express is that the final one should be 'def () -> cm.SecondSubClass'. Is there a neat way of expressing what I want?

[OddBloke]

from typing import ContextManager

class _SuperClass: ...

class SuperClass(ContextManager[SuperClass], _SuperClass): ...

class FirstSubClass(SuperClass): ...

class SecondSubClass(_SuperClass, ContextManager[SecondSubClass]): ...


reveal_type(SuperClass().__enter__)
reveal_type(FirstSubClass().__enter__)
reveal_type(SecondSubClass().__enter__)

gives

cm.py:13: error: Revealed type is 'def () -> cm.SuperClass*'
cm.py:14: error: Revealed type is 'def () -> cm.SuperClass*'
cm.py:15: error: Revealed type is 'def () -> cm.SecondSubClass*'

which is what I'm looking for.

However, this will cause problems if modifications happen in SuperClass instead of _SuperClass (as they will then only be inherited by FirstSubClass and not SecondSubClass).

[OddBloke]

(Note that I don't care if FirstSubClass retains SuperClass as its ContextManager type; I'd be happy with it either taking FirstSubClass or SuperClass.)

[ilevkivskyi]

In some sense mypy is actually right here, it just follows MRO and finds the first __enter__. But why don't you just want to add an explicit annotation to __enter__ in SecondSubClass? I assume that the method is different form __enter__ in the superclass since you want a different type for it. Otherwise you can use a self-type:

T = TypeVar('T', bound=SuperClass)
class SuperClass(ContextManager[SuperClass]):
    def __enter__(self: T) -> T: ...
class FirstSubClass(SuperClass): ...
class SecondSubClass(SuperClass, ContextManager[SecondSubClass]): ...

reveal_type(SuperClass().__enter__)  # Revealed type is 'def () -> test.SuperClass*'
reveal_type(FirstSubClass().__enter__)  #  Revealed type is 'def () -> test.FirstSubClass*'
reveal_type(SecondSubClass().__enter__)  # Revealed type is 'def () -> test.SecondSubClass'

Also note that ContextManager will probably be made a protocol at some point, so that it will be not necessary to add an explicit base class.

[OddBloke]

Yeah, I see why it's happening with the way my code is but was missing the pieces you've just given me to work out what else I could do. Thanks!