No type widening on generic self argument

[NiklasRosenstein]

Bug Report

Sorry for the bad title, I've tried to come up with something short and descriptive here.

The issue I am facing occurs when specializing the type of a self argument on a generic class.

To Reproduce

import typing as t

T = t.TypeVar('T')

def concat(its: t.Iterable[t.Iterable[T]]) -> t.Iterable[T]:
  for x in its:
    for y in x:
      yield y

v = ['abc', 'def']
reveal_type(v)                  # Revealed type is 'builtins.list[builtins.str*]'
print(''.join(concat(v)))       # works!


class Stream(t.Generic[T]):
  def __init__(self, it: t.Iterable[T]) -> None:
    self.it = iter(it)
  def __iter__(self) -> t.Iterator[T]:
    return self.it
  def concat(self: 'Stream[t.Iterable[T]]') -> 'Stream[T]':
    return Stream(concat(self))

s = Stream(v)
reveal_type(s)                  # Revealed type is 'test.Stream[builtins.str*]'
reveal_type(s.it)               # Revealed type is 'typing.Iterator[builtins.str*]'
print(''.join(concat(s)))       # works!
print(''.join(s.concat()))      # Invalid self argument "Stream[str]" to attribute function "concat" with type "Callable[[Stream[Iterable[T]]], Stream[T]]"

Expected Behavior

I expect the s.concat() call to be accepted by MyPy.

It actually works when you decorate the self argument with t.Iterable[t.Iterable[T]] instead.

  def concat(self: 't.Iterable[t.Iterable[T]]') -> 'Stream[T]':
    return Stream(concat(self))

Your Environment

• Mypy version used: 0.812
• Mypy command-line flags: mypy test.py
• Mypy configuration options from mypy.ini (and other config files): n/a
• Python version used: 3.8.4
• Operating system and version: WSL2 Ubuntu 20

[erictraut]

I think the current behavior is correct. The type of T has already been established as Iterable[str] when s was specialized. When you then try to call s.concat(), Stream[str] is not compatible with Stream[Iterable[str]] because T is invariant.

You could special-case str (since iterating over a str outputs other str instances) using an overload:

    @t.overload
    def concat(self: "Stream[str]") -> "Stream[str]":
        ...

    @t.overload
    def concat(self: "Stream[t.Iterable[T]]") -> "Stream[T]":
        ...

    def concat(self: "t.Union[Stream[t.Iterable[t.Any]], Stream[str]]") -> t.Any:
        return Stream(concat(self))

I verified that this type checks fine in both mypy and pyright.

[NiklasRosenstein]

Thanks @erictraut , that seems to work in this scenario.

I'm not 100% sure why I didn't notice it earlier, but it appears that for the normal case (ie. not Stream[str] but Stream[t.Iterable[T]]) it also doesn't quite work either.

import typing as t

T = t.TypeVar('T')

class Stream(t.Generic[T]):
  def __init__(self, it: t.Iterable[T]) -> None:
    self.it = iter(it)
  def __iter__(self) -> t.Iterator[T]:
    return self.it
  def concat(self: 'Stream[t.Iterable[T]]') -> 'Stream[T]':
    return Stream([])  # todo

s = Stream([[1, 2, 3], [4, 5]])
reveal_type(s)           # note: Revealed type is 't.Stream[builtins.list*[builtins.int]]'
reveal_type(s.concat())  # error: Invalid self argument "Stream[List[int]]" to attribute function "concat" with type "Callable[[Stream[Iterable[T]]], Stream[T]]"
                         # \note: Revealed type is 't.Stream[builtins.list*[builtins.int]]'

Is this not supposed to work or am I still missing a crucial element?

Replacing Stream[t.Iterable[T]] with t.Iterable[t.Iterable[T]] get's rid the error, but the revealed type is still t.Stream[builtins.list*[builtins.int]] rather than t.Stream[builtins.int].

[erictraut]

No, I wouldn't expect this to work. The problem is that s is type Stream[List[int]]. The Stream class is already specialized in this case, so it as if all instances of T within the class are replaced with List[int]. That means that the concat method is specialized to concat(self: Stream[Iterable[List[int]]]) -> Stream[List[int]]. The expression s.concat attempts to bind s to the self parameter, and there's a type mismatch because Stream[List[int]] (the type of s) is not compatible with Stream[Iterable[List[int]]] (the specialized type of self).

I'm not quite sure what you're intending to do with this code. Normally "concat" involves concatenating one stream to another, so perhaps you intended for concat to take an additional parameter? Also, it appears that Stream is already an Iterable (since it supports the __iter__ method), so I'm not sure why you are adding a second level of Iterable within the type argument for the self parameter. If you remove that (or, equivalently, remove the type annotation for self completely), it will type check fine.

[hauntsaninja]

In addition to what Eric mentioned, your TypeVar is invariant, so a Stream[List[int]] "is not a" Stream[Iterable[int]]. See https://mypy.readthedocs.io/en/stable/generics.html#variance-of-generic-types
I think the concat is meant to be a flatten kind of operation.

[erictraut]

Ah, that makes more sense if it's a "flatten" operation. Perhaps this addresses your needs?

import typing as t

T = t.TypeVar("T", covariant=True)

class Stream(t.Generic[T]):
    def __init__(self, it: t.Iterable[T]) -> None:
        self.it = iter(it)

    def __iter__(self) -> t.Iterator[T]:
        return self.it

    @staticmethod
    def flatten(s: "Stream[t.Iterable[T]]") -> "Stream[T]":
        ...

s = Stream([[1, 2, 3], [4, 5]])
reveal_type(s)
reveal_type(Stream.flatten(s))

[NiklasRosenstein]

Hey @erictraut ,

I'm back at this again 😄 Why does it have to be a staticmethod? Shouldn't this be considered a bug? In a way Stream.flatten(s) isn't really different from s.flatten() (except for the dynamic method lookup),

This behavior definitely seems odd to me (note how Mypy accepts it without @staticmethod as well).

import typing as t

T = t.TypeVar("T", covariant=True)

class Stream(t.Generic[T]):
    def __init__(self, it: t.Iterable[T]) -> None:
        self.it = iter(it)

    def __iter__(self) -> t.Iterator[T]:
        return self.it

    def flatten(s: "Stream[t.Iterable[T]]") -> "Stream[T]":
        ...

s = Stream([[1, 2, 3], [4, 5]])
reveal_type(s)                     # Revealed type is "test.Stream[builtins.list*[builtins.int]]"
reveal_type(Stream.flatten(s))     # Revealed type is "test.Stream[builtins.int*]"
reveal_type(s.flatten())           # Revealed type is "test.Stream[builtins.list*[builtins.int]]"

[erictraut]

In your sample above, you're invoking an instance method on an object with a specialized type. The type of object s is Stream[list[int]]. In the specialized version of the class, T is effectively replaced with list[int] everywhere in the implementation. The type of flatten therefore becomes (self: Stream[Iterable[list[int]]) -> Stream[list[int]], but when flatten is bound to the object, the binding operation attempts to assign type Stream[list[int]] to parameter self. Since Stream[list[int]] is incompatible with Stream[Iterable[list[int]], the binding will fail. I'm not sure why mypy allows this without generating an error. Pyright generates an error in this case.

Here's a simpler example that might make it clearer what I mean. Mypy does generate an error in this case.

T = TypeVar("T")

class A(Generic[T]):
    def foo(self: "A[int]"):
        ...

A[int]().foo()  # No error
A[str]().foo()  # Error: Could not bind method "foo" because "A[str]" is not assignable to parameter "self"

An instance method will work here if you use a different type variable in that method — one that is scoped to flatten, not to the class.

T = t.TypeVar("T", covariant=True)
U = t.TypeVar("U", covariant=True)

class Stream(t.Generic[T]):
    def __init__(self, it: t.Iterable[T]) -> None:
        self.it = iter(it)

    def __iter__(self) -> t.Iterator[T]:
        return self.it

    def flatten(self: "Stream[t.Iterable[U]]") -> "Stream[U]":
        ...