How to type __iter__ with a TypeVar?

Question

I am trying to make a dataclass with a generic type that I can unpack and feed as an argument to numpy's linspace. For this I need to give a return type for __iter__ with a TypeVar:

from typing import Iterator, Union, Generic, TypeVar, Any
import numpy as np

VT = TypeVar("VT", bound=float)

class Arguments(Generic[VT]):
    def __init__(self, start: VT, stop: VT, num: int):
        self.start = start
        self.stop = stop
        self.num = num

    def __iter__(self) -> Iterator[Union[VT, int]]:
        return iter([self.start, self.stop, self.num])

args: Arguments[float] = Arguments(1.2, 2.5, 10)

print(np.linspace(*args))

The execution works fine, but mypy (0.920) fails with the following error:

$ mypy test.py
test.py:17: error: No overload variant of "linspace" matches argument type "Arguments[float]"  [call-overload]
test.py:17: note: Possible overload variants:
test.py:17: note:     def linspace(start: Union[Any, Sequence[Any], ndarray[Any, Any], Any, Sequence[Sequence[Any]]], stop: Union[Any, Sequence[Any], ndarray[Any, Any], Any, Sequence[Sequence[Any]]], num: SupportsIndex = ..., endp
oint: bool = ..., retstep: Literal[False] = ..., dtype: Union[dtype[Any], None, type, _SupportsDType[dtype[Any]], str, Union[Tuple[Any, int], Tuple[Any, Union[SupportsIndex, Sequence[SupportsIndex]]], List[Any], _DTypeDict, Tuple[
Any, Any]]] = ..., axis: SupportsIndex = ...) -> ndarray[Any, Any]
test.py:17: note:     def linspace(start: Union[Any, Sequence[Any], ndarray[Any, Any], Any, Sequence[Sequence[Any]]], stop: Union[Any, Sequence[Any], ndarray[Any, Any], Any, Sequence[Sequence[Any]]], num: SupportsIndex = ..., endp
oint: bool = ..., retstep: Literal[True] = ..., dtype: Union[dtype[Any], None, type, _SupportsDType[dtype[Any]], str, Union[Tuple[Any, int], Tuple[Any, Union[SupportsIndex, Sequence[SupportsIndex]]], List[Any], _DTypeDict, Tuple[A
ny, Any]]] = ..., axis: SupportsIndex = ...) -> Tuple[ndarray[Any, Any], Any]
Found 1 error in 1 file (checked 1 source file)

If I do not make the class generic, mypy does not find issues:

from typing import Iterator, Union
import numpy as np

class Arguments:
    def __init__(self, start: float, stop: float, num: int):
        self.start = start
        self.stop = stop
        self.num = num

    def __iter__(self) -> Iterator[Union[float, int]]:
        return iter([self.start, self.stop, self.num])

args: Arguments = Arguments(1.2, 2.5, 10)

print(np.linspace(*args))

If I use a simpler linspace, I get a different error message:

from typing import Iterator, Union, Generic, TypeVar, Any

VT = TypeVar("VT", bound=float)

class Arguments(Generic[VT]):
    def __init__(self, start: VT, stop: VT, num: int):
        self.start = start
        self.stop = stop
        self.num = num

    def __iter__(self) -> Iterator[Union[VT, int]]:
        return iter([self.start, self.stop, self.num])

args: Arguments[float] = Arguments(1.2, 2.5, 10)

def my_linspace(start: float, stop: float, num: int) -> Any:
    return start

print(my_linspace(*args))

gives

$ mypy test.py
test.py:19: error: Argument 1 to "my_linspace" has incompatible type "*Arguments[float]"; expected "int"  [arg-type]
Found 1 error in 1 file (checked 1 source file)

What should I use as a return type to __iter__ instead of Iterator[Union[VT, int]]?

EDIT: Following the release of mypy 0.930 or 0.931, the above code works fine.

user_4111130 · Answer

A solution can be to use a non-generic subclass of Tuple, like this:

def my_linspace(start: float, stop: float, num: int) -> Any:
    return start


class Arguments(Tuple[float, float, int]):

    def __new__(cls, *args):
        return super().__new__(cls, args)

    def __init__(self, start: float, stop: float, num: int):
        super().__init__()


y: Arguments = Arguments(1.0, 2.0, 3)
print(my_linspace(*y))

How to type iter with a TypeVar?

1 Answers