TypeVar incompatible with constrained union

[NeilGirdhar]

from typing import Optional, TypeVar

U = TypeVar('U', bound=Optional[int])


def f(u: U) -> U:
 if u is None:
 return None
 assert isinstance(u, int) # removing this line causes it to pass.
 return u

Gives

main.py:10: error: Incompatible return value type (got "int", expected "U")

I cannot remove the assertion since it stands in for other code that is effectively constraining u.

This can be expressed tediously using overload:

from typing import Optional, TypeVar, overload


@overload
def f(u: None) -> None: ...

@overload
def f(u: int) -> int: ...

def f(u: Optional[int]) -> Optional[int]:
 if u is None:
 return None
 assert isinstance(u, int)
 return u

Would it be possible to get the TypeVar approach to pass? It is far more succinct, and this quickly blows up in complexity if there are more argument-return type constraints to be enforced.

This pattern is analogous to the common class factory pattern:

class C: pass
class B(C): pass
class A(B): pass

T = TypeVar('T', bound=C)
def factory(cls: Type[T]) -> T:
 return cls()

Since A < B < C, this enforces

factory(Type[C]) -> C
factory(Type[B]) -> B
factory(Type[A]) -> A

This issue argues that since int < Optional[int] and None < Optional[int], we should ideally be able to similarly specify that

f(int) -> int
f(None) -> None

using the same notation.

[habnabit]

TypeVar isn't how you create aliases:

from typing import Optional

U = Optional[int]

def f(u: U) -> U:
 if u is None:
 return None
 assert isinstance(u, int)
 return u

.. though U isn't a great name for that.

[NeilGirdhar]

@habnabit I'm not trying to create an alias, I'm trying to constrain the return type to match the argument type.

[NeilGirdhar]

Related to #9003 and duplicate of #5720.

[hauntsaninja]

You can get your code snippet to pass using a value restricted typevar, instead of a bound:

from typing import Optional, TypeVar

U = TypeVar('U', None, int)

def f(u: U) -> U:
 if u is None:
 return None
 reveal_type(u) # reveals an int
 return u

Note a problematic case with bounded typevars to be mindful of in situations similar to this:

U = TypeVar('U', bound=Optional[int])

def f(u: U) -> U:
 if u is None:
 return None
 assert isinstance(u, int)
 return 0 # if this were to typecheck, it would be unsound if a subclass of int is passed in

(edit: clarifying that bounds aren't central to the issue here, eg #8354 (comment))

[gvanrossum]

Interesting. For that snippet, mypy complains about return 0, pyre complains about both return lines, pyright and pytype don't complain about either.

[NeilGirdhar]

@hauntsaninja Could you give some insight into the difference between TypeVar('U', None, int) and TypeVar('U', bound=Optional[int])? It might not always be possible to break up a union type when the bound type is coming from a type alias. Also, the value restriction breaks callers:

from typing import Optional, TypeVar

U = TypeVar('U', None, int)

def f(u: U) -> U:
 if u is None:
 return None
 assert isinstance(u, int)
 return u
 
def g(v: Optional[int]) -> Optional[int]:
 return f(v) # error: Value of type variable "U" of "f" cannot be "Optional[int]"

@gvanrossum Looks like MyPy is the only one doing the right thing? That's great!

[hauntsaninja]

The relevant difference here is how subclasses would get treated. The following links have more info:
https://mypy.readthedocs.io/en/stable/generics.html#type-variables-with-value-restriction
https://www.python.org/dev/peps/pep-0484/#generics

Callers won't break if they pass in int, None or a similar typevar type. That said, I think your example above is something mypy could support, eg, since the following type checks:

from typing import Optional, TypeVar

U = TypeVar('U', None, int)

def f(u: U) -> U:
 if u is None:
 return None
 return u
 
def g(v: Optional[int]) -> Optional[int]:
 if isinstance(v, int):
 return f(v)
 return f(v)

Anyway, for now I think you're probably still best served by overloads or using a typevar and type-ignoring within the function. Note if you type ignore, it's really easy to forgo type safety in the presence of int subclasses, eg, even the following wouldn't be type safe:

from typing import Optional, TypeVar

U = TypeVar('U', bound=Optional[int])

def f(u: U) -> U:
 if u is None:
 return None
 return u + 0

But yes, as long as we're careful, I can't see a reason why mypy shouldn't accept:

from typing import Optional, TypeVar

U = TypeVar('U', bound=Optional[int])

def sleep(x: int) -> None: ...

def f(u: U) -> U:
 if u is None:
 return None
 sleep(u)
 return u

It should probably also not think the return None is unreachable.

[NeilGirdhar]

@hauntsaninja Thanks, that makes perfect sense.

Just a guess, but it seems that in:

def f(u: U) -> U:
 if u is None:
 return None
 assert isinstance(u, int)
 return u

at return u, mypy has forgotten that u is constrained to be of type U, and replaced that constraint with type int. In fact, u has type Intersection[U, int]. So perhaps this will be easier to implement if an Intersection type is added.

Thanks for the explanation.

[florensacc]

This problem also happens with bound that are not unions:

class A: pass
class B(A): pass
T = TypeVar("T", bound=A)

def foo(t: T) -> T:
 if isinstance(t, B):
 return t # this errors
 return t # this is fine 

This will return Incompatible return value type (got "B", expected "T")

[gvanrossum]

But what if we pass a C and expect a C back?

[florensacc]

Well, in my example you're returning the exact same object t, so no matter what class it is, you'll get the same class "back".
For more complex examples where you don't return the exact same object, like this example I found with

def foo(t: T) -> T:
 if isinstance(t, B):
 return B()

I agree that this should error because passing any subclass of B would not return an object of exactly the same class.

But in general, shouldn't it be possible to say "whatever the input class is, the output is of EXACTLY the same class"?

Does this comment address your "what if we pass a C" question @gvanrossum ?

[gvanrossum]

In your first example, after isinstance(t, B), the type of t is "narrowed" to B, so mypy treats it the same (from the POV of types) as your second example. The second example shows that B is not compatible with T, so it is not correct, and therefore neither is the first example. (The type checker doesn't understand enough of the data/control flow through your code to realize that in the first example we actually return the same object.)

In any case, your example is not the same as the issue being discussed in this issue. If you need more help please consult with the helpful folks on Gitter first.