Add basic support for polymorphic infernce with ParamSpec

python · ilevkivskyi · Aug 15, 2023 · Aug 7, 2023 · Aug 8, 2023 · Aug 8, 2023
commit 9b4eac81eed87577388474559a00bd72ed09bd80
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -122,6 +122,7 @@
     false_only,
     fixup_partial_type,
     function_type,
+    get_all_type_vars,
     get_type_vars,
     is_literal_type_like,
     make_simplified_union,
@@ -145,6 +146,7 @@
     LiteralValue,
     NoneType,
     Overloaded,
+    Parameters,
     ParamSpecFlavor,
     ParamSpecType,
     PartialType,
@@ -167,6 +169,7 @@
     get_proper_types,
     has_recursive_types,
     is_named_instance,
+    remove_dups,
     split_with_prefix_and_suffix,
 )
 from mypy.types_utils import (
@@ -5632,18 +5635,24 @@ def __init__(self, poly_tvars: Sequence[TypeVarLikeType]) -> None:
         self.bound_tvars: set[TypeVarLikeType] = set()
         self.seen_aliases: set[TypeInfo] = set()
 
-    def visit_callable_type(self, t: CallableType) -> Type:
-        found_vars = set()
+    def collect_vars(self, t: CallableType | Parameters) -> list[TypeVarLikeType]:
+        found_vars = []
         for arg in t.arg_types:
-            found_vars |= set(get_type_vars(arg)) & self.poly_tvars
+            found_vars += [
+                tv
+                for tv in get_all_type_vars(arg)
+                if tv in self.poly_tvars and tv not in self.bound_tvars
+            ]
+        return remove_dups(found_vars)
 
-        found_vars -= self.bound_tvars
-        self.bound_tvars |= found_vars
+    def visit_callable_type(self, t: CallableType) -> Type:
+        found_vars = self.collect_vars(t)
+        self.bound_tvars |= set(found_vars)
         result = super().visit_callable_type(t)
-        self.bound_tvars -= found_vars
+        self.bound_tvars -= set(found_vars)
 
         assert isinstance(result, ProperType) and isinstance(result, CallableType)
-        result.variables = list(result.variables) + list(found_vars)
+        result.variables = list(result.variables) + found_vars
         return result
 
     def visit_type_var(self, t: TypeVarType) -> Type:
@@ -5652,8 +5661,9 @@ def visit_type_var(self, t: TypeVarType) -> Type:
         return super().visit_type_var(t)
 
     def visit_param_spec(self, t: ParamSpecType) -> Type:
-        # TODO: Support polymorphic apply for ParamSpec.
-        raise PolyTranslationError()
+        if t in self.poly_tvars and t not in self.bound_tvars:
+            raise PolyTranslationError()
+        return super().visit_param_spec(t)
 
     def visit_type_var_tuple(self, t: TypeVarTupleType) -> Type:
         # TODO: Support polymorphic apply for TypeVarTuple.
@@ -5669,6 +5679,21 @@ def visit_type_alias_type(self, t: TypeAliasType) -> Type:
         raise PolyTranslationError()
 
     def visit_instance(self, t: Instance) -> Type:
+        if t.type.has_param_spec_type:
+            param_spec_index = next(
+                i for (i, tv) in enumerate(t.type.defn.type_vars) if isinstance(tv, ParamSpecType)
+            )
+            p = get_proper_type(t.args[param_spec_index])
+            if isinstance(p, Parameters):
+                found_vars = self.collect_vars(p)
+                self.bound_tvars |= set(found_vars)
+                new_args = [a.accept(self) for a in t.args]
+                self.bound_tvars -= set(found_vars)
+
+                repl = new_args[param_spec_index]
+                assert isinstance(repl, ProperType) and isinstance(repl, Parameters)
+                repl.variables = list(repl.variables) + list(found_vars)
+                return t.copy_modified(args=new_args)
         # There is the same problem with callback protocols as with aliases
         # (callback protocols are essentially more flexible aliases to callables).
         # Note: consider supporting bindings in instances, e.g. LRUCache[[x: T], T].

diff --git a/mypy/constraints.py b/mypy/constraints.py
@@ -899,7 +899,7 @@ def visit_callable_type(self, template: CallableType) -> list[Constraint]:
             cactual = self.actual.with_unpacked_kwargs()
             param_spec = template.param_spec()
             if param_spec is None:
-                # FIX verify argument counts
+                # TODO: verify argument counts; more generally, use some "formal to actual" map
                 # TODO: Erase template variables if it is generic?
                 if (
                     type_state.infer_polymorphic
@@ -943,34 +943,52 @@ def visit_callable_type(self, template: CallableType) -> list[Constraint]:
                         cactual_args = cactual.arg_types
                     # The lengths should match, but don't crash (it will error elsewhere).
                     for t, a in zip(template_args, cactual_args):
+                        if isinstance(a, ParamSpecType) and not isinstance(t, ParamSpecType):
+                            # This avoids bogus constraints like T <: P.args
+                            # TODO: figure out a more principled way to skip arg_kind mismatch
+                            # (see also a similar to do item in corresponding branch below)
+                            continue
                         # Negate direction due to function argument type contravariance.
                         res.extend(infer_constraints(t, a, neg_op(self.direction)))
             else:
                 # sometimes, it appears we try to get constraints between two paramspec callables?
 
-                # TODO: Direction
                 # TODO: check the prefixes match
                 prefix = param_spec.prefix
                 prefix_len = len(prefix.arg_types)
                 cactual_ps = cactual.param_spec()
 
+                if type_state.infer_polymorphic and cactual.variables and not self.skip_neg_op:
+                    # Similar logic to the branch above.
+                    res.extend(
+                        infer_constraints(
+                            cactual, template, neg_op(self.direction), skip_neg_op=True
+                        )
+                    )
+                    extra_tvars = True
+
                 if not cactual_ps:
                     max_prefix_len = len([k for k in cactual.arg_kinds if k in (ARG_POS, ARG_OPT)])
                     prefix_len = min(prefix_len, max_prefix_len)
                     res.append(
                         Constraint(
                             param_spec,
-                            SUBTYPE_OF,
-                            cactual.copy_modified(
+                            neg_op(self.direction),
+                            Parameters(
                                 arg_types=cactual.arg_types[prefix_len:],
                                 arg_kinds=cactual.arg_kinds[prefix_len:],
                                 arg_names=cactual.arg_names[prefix_len:],
-                                ret_type=UninhabitedType(),
+                                variables=cactual.variables
+                                if not type_state.infer_polymorphic
+                                else [],
                             ),
                         )
                     )
                 else:
-                    res.append(Constraint(param_spec, SUBTYPE_OF, cactual_ps))
+                    if not param_spec.prefix.arg_types or cactual_ps.prefix.arg_types:
+                        # TODO: figure out a more general logic to reject shorter prefix in actual.
+                        # This may be actually fixed by a more general to do item above.
+                        res.append(Constraint(param_spec, neg_op(self.direction), cactual_ps))
 
                 # compare prefixes
                 cactual_prefix = cactual.copy_modified(
@@ -979,7 +997,6 @@ def visit_callable_type(self, template: CallableType) -> list[Constraint]:
                     arg_names=cactual.arg_names[:prefix_len],
                 )
 
-                # TODO: see above "FIX" comments for param_spec is None case
                 # TODO: this assumes positional arguments
                 for t, a in zip(prefix.arg_types, cactual_prefix.arg_types):
                     res.extend(infer_constraints(t, a, neg_op(self.direction)))

diff --git a/mypy/solve.py b/mypy/solve.py
@@ -12,7 +12,7 @@
 from mypy.join import join_types
 from mypy.meet import meet_type_list, meet_types
 from mypy.subtypes import is_subtype
-from mypy.typeops import get_type_vars
+from mypy.typeops import get_all_type_vars
 from mypy.types import (
     AnyType,
     Instance,
@@ -463,4 +463,4 @@ def check_linear(scc: set[TypeVarId], lowers: Bounds, uppers: Bounds) -> bool:
 
 def get_vars(target: Type, vars: list[TypeVarId]) -> set[TypeVarId]:
     """Find type variables for which we are solving in a target type."""
-    return {tv.id for tv in get_type_vars(target)} & set(vars)
+    return {tv.id for tv in get_all_type_vars(target)} & set(vars)
diff --git a/mypy/typeops.py b/mypy/typeops.py
@@ -952,22 +952,31 @@ def coerce_to_literal(typ: Type) -> Type:
 
 
 def get_type_vars(tp: Type) -> list[TypeVarType]:
-    return tp.accept(TypeVarExtractor())
+    return cast("list[TypeVarType]", tp.accept(TypeVarExtractor()))
 
 
-class TypeVarExtractor(TypeQuery[List[TypeVarType]]):
-    def __init__(self) -> None:
+def get_all_type_vars(tp: Type) -> list[TypeVarLikeType]:
+    # TODO: should we always use this function instead of get_type_vars() above?
+    return tp.accept(TypeVarExtractor(include_all=True))
+
+
+class TypeVarExtractor(TypeQuery[List[TypeVarLikeType]]):
+    def __init__(self, include_all: bool = False) -> None:
         super().__init__(self._merge)
+        self.include_all = include_all
 
-    def _merge(self, iter: Iterable[list[TypeVarType]]) -> list[TypeVarType]:
+    def _merge(self, iter: Iterable[list[TypeVarLikeType]]) -> list[TypeVarLikeType]:
         out = []
         for item in iter:
             out.extend(item)
         return out
 
-    def visit_type_var(self, t: TypeVarType) -> list[TypeVarType]:
+    def visit_type_var(self, t: TypeVarType) -> list[TypeVarLikeType]:
         return [t]
 
+    def visit_param_spec(self, t: ParamSpecType) -> list[TypeVarLikeType]:
+        return [t] if self.include_all else []
+
 
 def custom_special_method(typ: Type, name: str, check_all: bool = False) -> bool:
     """Does this type have a custom special method such as __format__() or __eq__()?

diff --git a/test-data/unit/check-generics.test b/test-data/unit/check-generics.test
@@ -3035,3 +3035,39 @@ reveal_type(dec1(id2))  # N: Revealed type is "def [S in (builtins.int, builtins
 reveal_type(dec2(id1))  # N: Revealed type is "def [UC <: __main__.C] (UC`5) -> builtins.list[UC`5]"
 reveal_type(dec2(id2))  # N: Revealed type is "def (<nothing>) -> builtins.list[<nothing>]" \
                         # E: Argument 1 to "dec2" has incompatible type "Callable[[V], V]"; expected "Callable[[<nothing>], <nothing>]"
+
+[case testInferenceAgainstGenericParamSpecBasicInList]
+# flags: --new-type-inference
+from typing import TypeVar, Callable, List, Tuple
+from typing_extensions import ParamSpec
+
+T = TypeVar('T')
+P = ParamSpec('P')
+U = TypeVar('U')
+V = TypeVar('V')
+
+def dec(f: Callable[P, T]) -> Callable[P, List[T]]: ...
+def id(x: U) -> U: ...
+def either(x: U, y: U) -> U: ...
+def pair(x: U, y: V) -> Tuple[U, V]: ...
+reveal_type(dec(id))  # N: Revealed type is "def [T] (x: T`2) -> builtins.list[T`2]"
+reveal_type(dec(either))  # N: Revealed type is "def [T] (x: T`4, y: T`4) -> builtins.list[T`4]"
+reveal_type(dec(pair))  # N: Revealed type is "def [U, V] (x: U`-1, y: V`-2) -> builtins.list[Tuple[U`-1, V`-2]]"
+[builtins fixtures/list.pyi]
+
+[case testInferenceAgainstGenericParamSpecBasicDeList]
+# flags: --new-type-inference
+from typing import TypeVar, Callable, List, Tuple
+from typing_extensions import ParamSpec
+
+T = TypeVar('T')
+P = ParamSpec('P')
+U = TypeVar('U')
+V = TypeVar('V')
+
+def dec(f: Callable[P, List[T]]) -> Callable[P, T]: ...
+def id(x: U) -> U: ...
+def either(x: U, y: U) -> U: ...
+reveal_type(dec(id))  # N: Revealed type is "def [T] (x: builtins.list[T`2]) -> T`2"
+reveal_type(dec(either))  # N: Revealed type is "def [T] (x: builtins.list[T`4], y: builtins.list[T`4]) -> T`4"
+[builtins fixtures/list.pyi]
diff --git a/test-data/unit/check-overloading.test b/test-data/unit/check-overloading.test
@@ -6456,7 +6456,7 @@ P = ParamSpec("P")
 R = TypeVar("R")
 
 @overload
-def func(x: Callable[Concatenate[Any, P], R]) -> Callable[P, R]: ...  # E: Overloaded function signatures 1 and 2 overlap with incompatible return types
+def func(x: Callable[Concatenate[Any, P], R]) -> Callable[P, R]: ...
 @overload
 def func(x: Callable[P, R]) -> Callable[Concatenate[str, P], R]: ...
 def func(x: Callable[..., R]) -> Callable[..., R]: ...
@@ -6474,7 +6474,7 @@ eggs = lambda: 'eggs'
 reveal_type(func(eggs))  # N: Revealed type is "def (builtins.str) -> builtins.str"
 
 spam: Callable[..., str] = lambda x, y: 'baz'
-reveal_type(func(spam))  # N: Revealed type is "def (*Any, **Any) -> builtins.str"
+reveal_type(func(spam))  # N: Revealed type is "def (*Any, **Any) -> Any"
 
 [builtins fixtures/paramspec.pyi]
 

diff --git a/test-data/unit/check-parameter-specification.test b/test-data/unit/check-parameter-specification.test
@@ -1048,10 +1048,10 @@ class Job(Generic[_P, _T]):
 def generic_f(x: _T) -> _T: ...
 
 j = Job(generic_f)
-reveal_type(j)  # N: Revealed type is "__main__.Job[[x: _T`-1], _T`-1]"
+reveal_type(j)  # N: Revealed type is "__main__.Job[[x: _T`2], _T`2]"
 
 jf = j.into_callable()
-reveal_type(jf)  # N: Revealed type is "def [_T] (x: _T`-1) -> _T`-1"
+reveal_type(jf)  # N: Revealed type is "def [_T] (x: _T`2) -> _T`2"
 reveal_type(jf(1))  # N: Revealed type is "builtins.int"
 [builtins fixtures/paramspec.pyi]
 
@@ -1307,15 +1307,18 @@ reveal_type(bar(C(fn=foo, x=1)))  # N: Revealed type is "__main__.C[[x: builtins
 [builtins fixtures/paramspec.pyi]
 
 [case testParamSpecClassConstructor]
-from typing import ParamSpec, Callable
+from typing import ParamSpec, Callable, TypeVar
 
 P = ParamSpec("P")
 
 class SomeClass:
     def __init__(self, a: str) -> None:
         pass
 
-def func(t: Callable[P, SomeClass], val: Callable[P, SomeClass]) -> None:
+def func(t: Callable[P, SomeClass], val: Callable[P, SomeClass]) -> Callable[P, SomeClass]:
+    pass
+
+def func_regular(t: Callable[[T], SomeClass], val: Callable[[T], SomeClass]) -> Callable[[T], SomeClass]:
     pass
 
 def constructor(a: str) -> SomeClass:
@@ -1324,9 +1327,13 @@ def constructor(a: str) -> SomeClass:
 def wrong_constructor(a: bool) -> SomeClass:
     return SomeClass("a")
 
+def wrong_name_constructor(b: bool) -> SomeClass:
+    return SomeClass("a")
+
 func(SomeClass, constructor)
-func(SomeClass, wrong_constructor)  # E: Argument 1 to "func" has incompatible type "Type[SomeClass]"; expected "Callable[[VarArg(<nothing>), KwArg(<nothing>)], SomeClass]" \
-				    # E: Argument 2 to "func" has incompatible type "Callable[[bool], SomeClass]"; expected "Callable[[VarArg(<nothing>), KwArg(<nothing>)], SomeClass]"
+reveal_type(func(SomeClass, wrong_constructor))  # N: Revealed type is "def (a: <nothing>) -> __main__.SomeClass"
+reveal_type(func_regular(SomeClass, wrong_constructor))  # N: Revealed type is "def (<nothing>) -> __main__.SomeClass"
+func(SomeClass, wrong_name_constructor)  # E: Argument 1 to "func" has incompatible type "Type[SomeClass]"; expected "Callable[[<nothing>], SomeClass]"
 [builtins fixtures/paramspec.pyi]
 
 [case testParamSpecInTypeAliasBasic]
@@ -1547,5 +1554,5 @@ U = TypeVar("U")
 def dec(f: Callable[P, T]) -> Callable[P, List[T]]: ...
 def test(x: U) -> U: ...
 reveal_type(dec)  # N: Revealed type is "def [P, T] (f: def (*P.args, **P.kwargs) -> T`-2) -> def (*P.args, **P.kwargs) -> builtins.list[T`-2]"
-reveal_type(dec(test))  # N: Revealed type is "def [U] (x: U`-1) -> builtins.list[U`-1]"
+reveal_type(dec(test))  # N: Revealed type is "def [T] (x: T`2) -> builtins.list[T`2]"
 [builtins fixtures/paramspec.pyi]