Allow TypeVars in Concatenate

python · JukkaL · Apr 7, 2022 · Dec 24, 2021 · Dec 25, 2021 · Dec 26, 2021
commit 0e2b207b19bebe7a76f84baa015d133a8558256d
diff --git a/mypy/constraints.py b/mypy/constraints.py
@@ -461,7 +461,7 @@ def visit_instance(self, template: Instance) -> List[Constraint]:
                         if isinstance(suffix, Parameters):
                             # no such thing as variance for ParamSpecs
                             # TODO: is there a case I am missing?
-                            # TODO: what is setting meta_level to 0?
+                            # TODO: constraints between prefixes
                             prefix = mapped_arg.prefix
                             suffix = suffix.copy_modified(suffix.arg_types[len(prefix.arg_types):],
                                                           suffix.arg_kinds[len(prefix.arg_kinds):],
@@ -490,7 +490,7 @@ def visit_instance(self, template: Instance) -> List[Constraint]:
                         if isinstance(suffix, Parameters):
                             # no such thing as variance for ParamSpecs
                             # TODO: is there a case I am missing?
-                            # TODO: what is setting meta_level to 0?
+                            # TODO: constraints between prefixes
                             prefix = template_arg.prefix
                             suffix = suffix.copy_modified(suffix.arg_types[len(prefix.arg_types):],
                                                           suffix.arg_kinds[len(prefix.arg_kinds):],
@@ -586,7 +586,6 @@ def visit_callable_type(self, template: CallableType) -> List[Constraint]:
                         res.extend(infer_constraints(t, a, neg_op(self.direction)))
             else:
                 # TODO: Direction
-                # TODO: Deal with arguments that come before param spec ones?
                 # TODO: check the prefixes match
                 prefix = param_spec.prefix
                 prefix_len = len(prefix.arg_types)
@@ -597,6 +596,16 @@ def visit_callable_type(self, template: CallableType) -> List[Constraint]:
                                         arg_kinds=cactual.arg_kinds[prefix_len:],
                                         arg_names=cactual.arg_names[prefix_len:],
                                         ret_type=NoneType())))
+                # compare prefixes
+                cactual_prefix = cactual.copy_modified(
+                    arg_types=cactual.arg_types[:prefix_len],
+                    arg_kinds=cactual.arg_kinds[:prefix_len],
+                    arg_names=cactual.arg_names[:prefix_len])
+
+                #
8000
 TODO: see above "FIX" comments for param_spec is None case
+                # TODO: this assume positional arguments
+                for t, a in zip(prefix.arg_types, cactual_prefix.arg_types):
+                    res.extend(infer_constraints(t, a, neg_op(self.direction)))
 
             template_ret_type, cactual_ret_type = template.ret_type, cactual.ret_type
             if template.type_guard is not None:

diff --git a/mypy/expandtype.py b/mypy/expandtype.py
@@ -137,11 +137,14 @@ def visit_callable_type(self, t: CallableType) -> Type:
             # the replacement is ignored.
             if isinstance(repl, CallableType):
                 # Substitute *args: P.args, **kwargs: P.kwargs
-                t = t.expand_param_spec(repl)
-                # TODO: Substitute remaining arg types
-                return t.copy_modified(ret_type=t.ret_type.accept(self),
-                                       type_guard=(t.type_guard.accept(self)
-                                                   if t.type_guard is not None else None))
+                prefix = param_spec.prefix
+                t = t.expand_param_spec(repl, no_prefix=True)
+                return t.copy_modified(
+                    arg_types=self.expand_types(prefix.arg_types) + t.arg_types,
+                    arg_kinds=prefix.arg_kinds + t.arg_kinds,
+                    arg_names=prefix.arg_names + t.arg_names,
+                    ret_type=t.ret_type.accept(self),
+                    type_guard=(t.type_guard.accept(self) if t.type_guard is not None else None))
 
         return t.copy_modified(arg_types=self.expand_types(t.arg_types),
                                ret_type=t.ret_type.accept(self),

diff --git a/mypy/typeanal.py b/mypy/typeanal.py
@@ -847,10 +847,12 @@ def analyze_callable_args_for_concatenate(
         obj = self.named_type('builtins.object')
         # ick, CallableType should take ParamSpecType
         prefix = tvar_def.prefix
+        # we don't set the prefix here as generic arguments will get updated at some point
+        # in the future. CallableType.param_spec() accounts for this.
         return CallableType(
             [*prefix.arg_types,
              ParamSpecType(tvar_def.name, tvar_def.fullname, tvar_def.id, ParamSpecFlavor.ARGS,
-                           upper_bound=obj, prefix=tvar_def.prefix),
+                           upper_bound=obj),
              ParamSpecType(tvar_def.name, tvar_def.fullname, tvar_def.id, ParamSpecFlavor.KWARGS,
                            upper_bound=obj)],
             [*prefix.arg_kinds, nodes.ARG_STAR, nodes.ARG_STAR2],

diff --git a/mypy/types.py b/mypy/types.py
@@ -1520,10 +1520,15 @@ def param_spec(self) -> Optional[ParamSpecType]:
         return ParamSpecType(arg_type.name, arg_type.fullname, arg_type.id, ParamSpecFlavor.BARE,
                              arg_type.upper_bound, prefix=prefix)
 
-    def expand_param_spec(self, c: 'CallableType') -> 'CallableType':
-        return self.copy_modified(arg_types=self.arg_types[:-2] + c.arg_types,
-                                  arg_kinds=self.arg_kinds[:-2] + c.arg_kinds,
-                                  arg_names=self.arg_names[:-2] + c.arg_names)
+    def expand_param_spec(self, c: 'CallableType', no_prefix: bool = False) -> 'CallableType':
+        if no_prefix:
+            return self.copy_modified(arg_types=c.arg_types,
+                                      arg_kinds=c.arg_kinds,
+                                      arg_names=c.arg_names)
+        else:
+            return self.copy_modified(arg_types=self.arg_types[:-2] + c.arg_types,
+                                      arg_kinds=self.arg_kinds[:-2] + c.arg_kinds,
+                                      arg_names=self.arg_names[:-2] + c.arg_names)
 
     def __hash__(self) -> int:
         return hash((self.ret_type, self.is_type_obj(),

diff --git a/test-data/unit/check-parameter-specification.test b/test-data/unit/check-parameter-specification.test
@@ -582,3 +582,21 @@ def f2(c: Callable[P, R]) -> Callable[Concatenate[int, P], R]:
 # reason for rejection:
 f2(lambda x: 42)(42, x=42)
 [builtins fixtures/tuple.pyi]
+
+[case testParamSpecConcatenateWithTypeVar]
+from typing_extensions import ParamSpec, Concatenate
+from typing import Callable, TypeVar
+
+P = ParamSpec("P")
+R = TypeVar("R")
+S = TypeVar("S")
+
+def f(c: Callable[Concatenate[S, P], R]) -> Callable[Concatenate[S, P], R]: ...
+
+def a(n: int) -> None: ...
+
+n = f(a)
+
+reveal_type(n)  # N: Revealed type is "def (builtins.int*)"
+reveal_type(n(42))  # N: Revealed type is "None"
+[builtins fixtures/tuple.pyi]