Generic classes narrowed by @overloaded __init__

[srittau]

Quite often in typeshed we have the case that a generic class can be narrowed by the arguments provided to __init__. One example are file classes that are generic over AnyStr, where a mode argument narrows the type. It would be useful if this could be described somehow in typing, possibly by extending the @overload decorator or with some other mechanism.

[srittau]

One idea:

class MyFile(Generic[AnyStr]):
    @overload(narrow_generic_0=str)
    def __init__(self, mode: Literal["r", "w"] = ...) -> None: ...
    @overload(narrow_generic_0=bytes)
    def __init__(self, mode: Literal["rb", "wb", "b"]) -> None: ...
    @overload  # AnyStr is not narrowed
    def __init__(self, mode: str) -> None: ...

[JelleZijlstra]

I think this is doable by overloading __new__, like we implemented recently for subprocess.Popen.

[srittau]

That's a nice trick for stubs! I think something like this would be still be useful, though:

• For implementation files where the __new__ trick can't be used.
• As a more obvious/non-hacky alternative.
• To make stubs better reflect reality (by not including __new__).

[ilevkivskyi]

I think this is doable by overloading __new__, like we implemented recently for subprocess.Popen.

Yes, this is however rather a workaround than a real fix. This can cause troubles if an actual class uses __init__ at runtime. Recently there was a bug report caused by this, see python/typeshed#3363. Also allowing annotations on self and overloading in them would be useful for other things (e.g. supporting mixins and methods with signatures that depend on type argument). Also I don't think we need any special syntax/rules for this. We can just allow:

class MyFile(Generic[AnyStr]):
    @overload
    def __init__(self: MyFile[str], mode: Literal["r", "w"] = ...) -> None: ...
    @overload
    def __init__(self: MyFile[bytes], mode: Literal["rb", "wb", "b"]) -> None: ...
    @overload
    def __init__(self, mode: str) -> None: ...

From mypy point of view we would need to improve how bind_self() works plus few other things, I assume this would be similar for other type checkers.

UPDATE: mypy implementation may be actually more tricky and would need some special-casing in e.g. class_callable().

[ilevkivskyi]

Just to add an example for how annotations on self can be used to support mixins:

class Base:
    def base_meth(self) -> None: ...

class Mixin:
    def mix_meth(self: Base) -> None:
        # Here `self` is an intersection of current class and annotation.
        self.base_meth()  # OK

class Good(Base, Mixin): ...
class Bad(Mixin): ...

Good().mix_meth()  # OK
Bad().mix_meth()  # Error from `bind_self`

In principle we can make this work in mypy, but more advanced types for mixins may be available only if we add intersection types (at least for proper classes and their unions, similar to what is allowed in Type[...]).

[srittau]

One thing where this change is quite useful are my partial SQLalchemy stubs. Previously columns needed to be explicitly annotated:

class Foo(ORMBase):
    name: Column[str] = Column(String, primary_key=True)

Now using overloads this is no necessary anymore:

class Foo(ORMBase):
    name = Column(String, primary_key=True)

Needs six overloads per type and doesn't cover all corner cases, but still very useful.

[ilevkivskyi]

Have you thought about contributing whatever you have on SQLAlchemy to https://github.com/dropbox/sqlalchemy-stubs/? I think it would be great to have a single package with better coverage.

[ilevkivskyi]

@srittau ^^