Unexpected type binding when wrapping functions

[mailund]

Bug Report

I admit that I do not know if this is a bug, and not just my incomplete understanding of how the type system works, but I have run into some unexpected behaviour, and behaviour where mypy and pyre disagree, so it might be a bug.

I am writing code that lifts functions T -> R to Optional[T] -> Optional[R] in various ways, and when I write wrapper functions, type variables get bound in ways I didn't expect. It is particularly bad when I add protocols for various types on top of it, but I can create the unexpected behaviour without it. Basically, the issue is that if I write a function that changes a Callable[[T], R] into a Callable[[Optional[T]], Optional[R]], and then apply such a function to another generic function Callable[[T],R] the type variables are no longer generic, and I cannot apply the wrapped function. This happens both when I work with wrapper functions or when I implement the wrapper as a Generic class.

To Reproduce

If I write a wrapper with a function, it might look like this:

from typing import (
 TypeVar, 
 Protocol, Generic, 
 Optional as Opt
)


_T = TypeVar('_T')

class _F(Protocol[_T]):
 def __call__(self, __x: _T) -> _T:
 ...

# Lifting with a function...
def lift(f: _F[_T]) -> _F[Opt[_T]]:
 # This function would do more, but I just return f for the
 # type checking example
 return f # type: ignore

I get the desired behaviour if I wrap a function with a concrete type

def f(x: int) -> int:
 return x

reveal_type(lift(f)) # _F[Opt[int]] -- as expected

but if I use a generic type, the wrapped function is not generic:

def g(x: _T) -> _T:
 return x

reveal_type(lift(g)) # _F[Union[_T`-1,None]] -- _T is bound (incorrectly?)
 # pyre thinks this is _F[None] which is also odd...
lift(g)(1) # so won't accept int

Expected Behavior

I was expecting the wrapped function to have the type Callable[[Opt[_T]], Opt[_T]] and in particular I would expect to be able to call it with an int argument.

To Reproduce

If I use a Generic class for the lifted function, it can look like this:

# Moving generic type to Generic...
class lifted(Generic[_T]):
 """Lift type _T to Opt[_T]."""

 def __init__(self, f: _F[_T]) -> None:
 self._f = f

 def __call__(self, x: Opt[_T], /) -> Opt[_T]:
 """Call lifted."""
 return self._f(x) # type: ignore

reveal_type(lifted[_T]) # _F[T?] -> lifted[_T?]

and it will work for a function with a concrete type, like f: int -> int above

reveal_type(lifted(f)) # lifted[int] -- perfect
reveal_type(lifted(f).__call__) # Opt[int] -> Opt[int] -- perfect
lifted(f)(1) # We can call with int
lifted(f)(None) # We can call with None

but things get crazy when I use a generic function like g: _T -> _T:

reveal_type(lifted(g)) # lifted[_T`1] -- _T is bound!
 # pyre thinks (correctly I think) it is lifted[_T]
reveal_type(lifted(g).__call__) # Opt[_T`1] -> Opt[_T`1]
 # pyre says Any -> Any which I suppose is fine
lifted(g)(1) # incompatible type, int isn't an Opt[_T]
 # pyre seems to accep this one
lifted(g)(None) # but at least None is in Opt[_T]...

reveal_type(lifted[_T](g)) # lifted[_T?] -- Ok
reveal_type(lifted[_T](g).__call__) # def (None) ???
 # pyre: Opt[_T] -> Opt[_T] as expected

lifted[_T](g)(1) # incompatible type "int"; expected "None"
lifted[_T](g)(None) # but at least None works...
# Pyre doesn't like the [_T] here but accepts the calls without it...

Expected Behavior

Again, I would expect the Generic lifted class to accept types that match the wrapped function's type. This is more important when I need to add protocols to the types, but I would also expect it to work here.

Your Environment

I tested this in mypy playground and pyre playground

[erictraut]

Type variables are "solved" in the context of a call expression that targets the constructor of a generic class or a generic function. Only the type variables in the scope of that class or function are solved. Once a type variable goes out of scope, it can no longer be solved.

The expression listed(g) causes mypy to solve the TypeVar _T@lifted (i.e. _T in the scope of class lifted). The solution in this case is _T@g, which means that listed(g) evaluates to type lifted[_T@g]. However, at this point, _T@g is no longer in scope. It's an unsolved TypeVar from the scope of function g. It is not associated with the lifted class or its __call__ method, which means it cannot be solved in subsequent call expressions like lifted(g)(1).

FWIW, pyright largely agrees with mypy in this case.

So I don't think this is a bug in mypy. It's expected behavior.

[mailund]

Ok, I think I understand what you are saying, @erictraut, but let me just check. If I use a generic function or class with a type variable, I bind it when I apply it, and it is no longer functioning as a type variable. If I want the wrapped function to work, the wrapping has to know the concrete type.

So, for example, with a curry function (as a Generic)

from typing import (
 TypeVar, ParamSpec,
 Callable as Fn,
 Optional as Opt,
 Generic,
)

_T = TypeVar('_T')

class curry(Generic[_T]):
 def __init__(self, f: Fn[[_T, _T], _T]) -> None:
 self._f = f
 def __call__(self, x: _T) -> Fn[[_T], _T]:
 return lambda y: self._f(x, y)

I can't do something like

@curry
def f(x: _T, __y: _T) -> _T:
 return x

because _T isn't free after wrapping, and I cannot do this either

@curry[_T]
def f(x: _T, __y: _T) -> _T:
 return x

either, because I can't use a type variable like this when it is unbound?

I can, however, use it with concrete types

def f(x: _T, __y: _T) -> _T:
 return x

fi = curry[int](f)
reveal_type(fi) # curry[int]
reveal_type(fi(1)) # def int -> int
reveal_type(fi(1)(2)) # int

fs = curry[str](f)
reveal_type(fs) # curry[str]
reveal_type(fs("foo")) # str -> str
reveal_type(fs("foo")("bar")) # str

or I can use _T inside a function that uses the type variable because there it will be bound when I evaluate the function and thus get a concrete type there.

def bind(x: _T, f: Fn[[_T,_T], _T]) -> Fn[[_T], _T]:
 # ok here since _T is bound when we evaluate
 res = curry[_T](f)
 return res(x)
 
b = bind(12, f)
reveal_type(b) # int -> int

b2 = bind("foo", f)
reveal_type(b2) # str -> str

As long as I don't use an unbound type variable when instantiating the Generic here, it seems to work. Is that guaranteed, and is that the way to achieve what I tried above?

[sterliakov]

Just adding my 2 cents here, because I've faced this problem once before (and ended with very ugly pattern that violates DRY, see below).

Related: #1317

It is understandable why does mypy "solve" this variable earlier than we want. But there is no way to define the following decorator: "given function that takes fixed arguments set (say, f(x, y) that implements min signature, like f(x: _T, y: _T) -> _T), produce a function that has same signature with types Transformed[_T]", where Transformed is expressible with typing primitives and depends only on _T (propagating Optional is just one of possible use cases). It is important that _T is the same for decorated function and result, but we do not know _T value in advance.

I think that the following (very raw) approach should be better:

Given decorator:

from typing import Callable, SupportsInt, TypeVar

_T = TypeVar('_T')
_P = TypeVar('_P', bound=SupportsInt)

def allow_none(f: Callable[[_T], _T]) -> Callable[[_T | None], _T | None]:
 def inner(x: _T | None, /) -> _T | None:
 return None if x is None else f(x)
 return inner

@allow_none
def foo(x: _P) -> _P:
 if x is None:
 raise ValueError
 return x

After producing F[_P | None] type checker should recognize the fact that _P is used both on LHS and RHS of Callable definition. When it happens, _P should go back to F scope, allowing the expected substitution with respect to SupportsInt bound.

The same should happen is resulting type is Generic or Protocol: instead of producing generic class with unsolvable type variable, mypy should keep it unresolved. Now any function that performs similar operation (f(x: G[_T]) -> G[_T | None], where G is any Protocol or Generic) is malformed: when applied to generic with type variable G[_T], it produces G[_T | None] with _T out of scope, so only Any and None are compatible. It is definitely not the desired behavior.

I suppose that no existing code should be broken after such change. The resolution algorithm now produces unusable functions (they don't take expected arguments), so usage of such decorators must be very rare and complex. When decorated function is using concrete types, this approach will change nothing.

The only solution now is to explicitly re-declare type of decorated function. But why are they compatible, if foo(1) is OK while allow_none(_foo)(1) throws mypy error: E: Argument 1 has incompatible type "int"; expected "Optional[_P]"? I don't know, it seems to be inconsistent, LHS type is wider like x: int = 0.1, isn't it?

... # See example above

def _foo(x: _P) -> _P:
 if x is None:
 raise ValueError
 return x

foo: Callable[[_P | None], _P | None] = allow_none(_foo)

reveal_type(foo) # N: Revealed type is "def [_P <: typing.SupportsInt] (Union[_P`-1, None]) -> Union[_P`-1, None]"
reveal_type(foo(None)) # N: Revealed type is "None"
reveal_type(foo(1)) # N: Revealed type is "Union[builtins.int, None]"