__new__ type annotations have unexpected behavior in some cases

[itamarst]

Bug Report

I am trying to emulate some Pandas type behavior, where Index.__new__ can return different types. The simplified case I am trying to model is that sequences of np.datetime64 and datetime.datetime get turned into one class, and other objects get turned into another class.

The following is the only way I've gotten it to work, after trying many many variations:

# THIS VERSION WORKS, BUT REQUIRES MODIFYING CODE WITH EXTRA CLASS
from typing import TypeVar, Generic, List, Union, overload
from typing_extensions import Protocol
from datetime import datetime

T = TypeVar("T", covariant=True)
S = TypeVar("S")

class datetime64(int):
    """Stand-in for np.datetime64."""


class IndexType(Protocol[T]):
    def first(self) -> T: ...


class Index:

    @overload
    def __new__(cls, values: List[datetime64]) -> "Datetime64Index": ...
    @overload
    def __new__(cls, values: List[datetime]) -> "Datetime64Index": ...
    @overload
    def __new__(cls, values: List[S]) -> "DefaultIndex[S]": ...

    def __new__(cls, values):
        if type(values[0]) in (datetime, datetime64):
            cls = Datetime64Index
        else:
            cls = DefaultIndex
        return object.__new__(cls)


class DefaultIndex(Index, Generic[S]):
    def __init__(self, values: List[S]):
        self.values = values

    def first(self) -> S:
        return self.values[0]


class Datetime64Index(DefaultIndex):

    def __init__(self, values: Union[List[datetime], List[datetime64]]):
        self.values : List[datetime64] = [
            datetime64(o.timestamp()) if isinstance(o, datetime) else o
            for o in values
        ]

    def first(self) -> datetime:
        return datetime.fromtimestamp(self.values[0])


# Should work
a: IndexType[datetime] = Index([datetime64(100)])
b: IndexType[datetime] = Index([datetime(2000, 10, 20)])
c: IndexType[bool] = Index([True])

# Should complain
d: IndexType[datetime] = Index(["a"])
e: IndexType[bool] = Index(["a"])

As expected, mypy only complains about the last two lines:

$ mypy test.py
test.py:59: error: List item 0 has incompatible type "str"; expected "datetime64"
test.py:60: error: Incompatible types in assignment (expression has type "Datetime64Index", variable has type "IndexType[bool]")
test.py:60: note: Following member(s) of "Datetime64Index" have conflicts:
test.py:60: note:     Expected:
test.py:60: note:         def first(self) -> bool
test.py:60: note:     Got:
test.py:60: note:         def first(self) -> datetime
test.py:60: error: List item 0 has incompatible type "str"; expected "datetime64"

However, the need for DefaultIndex feels like a hack. What I would actually like to do is the following:

# THIS VERSION SHOULD WORK, BUT CAUSES MYPY TO ERRONEOUSLY(?) COMPLAIN
from typing import TypeVar, Generic, List, Union, overload
from typing_extensions import Protocol
from datetime import datetime

T = TypeVar("T", covariant=True)
S = TypeVar("S")

class datetime64(int):
    """Stand-in for np.datetime64."""


class IndexType(Protocol[T]):
    def first(self) -> T: ...


class Index(Generic[S]):

    @overload
    def __new__(cls, values: List[datetime64]) -> "Datetime64Index": ...
    @overload
    def __new__(cls, values: List[datetime]) -> "Datetime64Index": ...
    @overload
    def __new__(cls, values: List[S]) -> "Index[S]": ...

    def __new__(cls, values):
        if type(values[0]) in (datetime, datetime64):
            cls = Datetime64Index
        return object.__new__(cls)

    def __init__(self, values: List[S]):
        self.values = values

    def first(self) -> S:
        return self.values[0]


class Datetime64Index(Index):

    def __init__(self, values: Union[List[datetime], List[datetime64]]):
        self.values : List[datetime64] = [
            datetime64(o.timestamp()) if isinstance(o, datetime) else o
            for o in values
        ]

    def first(self) -> datetime:
        return datetime.fromtimestamp(self.values[0])


# Should work
a: IndexType[datetime] = Index([datetime64(100)])
b: IndexType[datetime] = Index([datetime(2000, 10, 20)])
c: IndexType[bool] = Index([True])

# Should complain
d: IndexType[datetime] = Index(["a"])
e: IndexType[bool] = Index(["a"])

However, mypy gets confused and complains about a: IndexType[datetime] = Index([datetime64(100)]):

test.py:50: error: List item 0 has incompatible type "datetime64"; expected "datetime"
test.py:55: error: List item 0 has incompatible type "str"; expected "datetime"
test.py:56: error: List item 0 has incompatible type "str"; expected "bool"
Found 3 errors in 1 file (checked 1 source file)

Your Environment

• Mypy version used: 0.782
• Mypy command-line flags: None
• Mypy configuration options from mypy.ini (and other config files): None
• Python version used: 3.7
• Operating system and version: Linux

[gvanrossum]

Generics are hard. :-(

It seems that there is interference between __init__ and __new__ and the type of Index([datetime64(100)]) is determined by looking at __init__, where it takes S to be datetime64.

Usually that's fine (normally __new__ and __init__ must match) but here it's not. :-(

Your working code works around the problem by separating __new__ and __init__ into different classes. I think that's the best you can do to model this API. It would seem that your toy example here could also be solved by making Index a function, but I presume there's something in the real API that precludes this.

Sorry, that's all I have -- actually fixing this would probably require some deep changes in mypy which I couldn't do myself.