Get basic Concatenate features working

This program: ```py from typing_extensions import ParamSpec, Concatenate from typing import Generic P = ParamSpec("P") class Z(Generic[P]): ... n: Z[[int]] reveal_type(n) def f(n: Z[P]) -> Z[Concatenate[bool, Concatenate[str, P]]]: ... reveal_type(f) reveal_type(f(n)) ``` outputs: ``` repro.py:10: note: Revealed type is "repro.Z[[builtins.int]]" repro.py:14: note: Revealed type is "def [P] (n: repro.Z[P`-1]) -> repro.Z[Concatenate[builtins.bool, builtins.str, P`-1]]" repro.py:15: note: Revealed type is "repro.Z[[builtins.bool, builtins.str, builtins.int]]" ``` Next up: checking inputs match prefixes and why does it only work when cache exists?
JukkaL · Apr 7, 2022 · Dec 24, 2021 · Dec 25, 2021 · Dec 26, 2021 · Dec 26, 2021
commit d9b352faa2d013259849952cb078a3b874f34932
diff --git a/mypy/expandtype.py b/mypy/expandtype.py
@@ -104,11 +104,21 @@ def visit_param_spec(self, t: ParamSpecType) -> Type:
         if isinstance(repl, Instance):
             inst = repl
             # Return copy of instance with type erasure flag on.
+            # TODO: what does prefix mean in this case?
+            # TODO: why does this case even happen? Instances aren't plural.
             return Instance(inst.type, inst.args, line=inst.line,
                             column=inst.column, erased=True)
         elif isinstance(repl, ParamSpecType):
-            return repl.with_flavor(t.flavor)
+            # TODO: what if both have prefixes???
+            # (realistically, `repl` is the unification variable for `t` so this is fine)
+            return repl.copy_modified(flavor=t.flavor, prefix=t.prefix)
+        elif isinstance(repl, Parameters):
+            return repl.copy_modified(t.prefix.arg_types + repl.arg_types,
+                                      t.prefix.arg_kinds + repl.arg_kinds,
+                                      t.prefix.arg_names + repl.arg_names)
         else:
+            # returning parameters
+            # TODO: should this branch be removed? better not to fail silently
             return repl
 
     def visit_parameters(self, t: Parameters) -> Type:

diff --git a/mypy/typeanal.py b/mypy/typeanal.py
@@ -266,6 +266,7 @@ def visit_unbound_type_nonoptional(self, t: UnboundType, defining_literal: bool)
                 return AnyType(TypeOfAny.special_form)
             # Concatenate is an operator, no need for a proper type
             elif node.fullname in ("typing_extensions.Concatenate", "typing.Concatenate"):
+                # TODO: detect valid locations (`allow_param_spec` is not here.)
                 return self.apply_concatenate_operator(t)
             else:
                 return self.analyze_unbound_type_without_type_info(t, sym, defining_literal)
@@ -280,12 +281,24 @@ def cannot_resolve_type(self, t: UnboundType) -> None:
             'Cannot resolve name "{}" (possible cyclic definition)'.format(t.name),
             t)
 
-    def apply_concatenate_operator(self, t: UnboundType) -> Optional[Type]:
+    def apply_concatenate_operator(self, t: UnboundType) -> Optional[ParamSpecType]:
         if len(t.args) == 0:
             self.api.fail('Concatenate needs type arguments', t)
             return AnyType(TypeOfAny.from_error)
 
-        raise RuntimeError("TODO")
+        # last argument has to be ParamSpec (or Concatenate)
+        ps = self.anal_type(t.args[-1], allow_param_spec=True)
+        if not isinstance(ps, ParamSpecType):
+            print(ps)
+            self.api.fail('The last parameter to Concatenate needs to be a ParamSpec', t)
+            return AnyType(TypeOfAny.from_error)
+
+        args = self.anal_array(t.args[:-1])
+        pre = ps.prefix
+        pre = Parameters(args + pre.arg_types,
+                         [ARG_POS] * len(args) + pre.arg_kinds,
+                         [None] * len(args) + pre.arg_names)
+        return ps.copy_modified(prefix=pre)
 
     def try_analyze_special_unbound_type(self, t: UnboundType, fullname: str) -> Optional[Type]:
         """Bind special type that is recognized through magic name such as 'typing.Any'.

diff --git a/mypy/types.py b/mypy/types.py
@@ -489,26 +489,43 @@ class ParamSpecType(TypeVarLikeType):
     always just 'object').
     """
 
-    __slots__ = ('flavor',)
+    __slots__ = ('flavor', 'prefix')
 
     flavor: int
+    prefix: 'Parameters'
 
     def __init__(
          self, name: str, fullname: str, id: Union[TypeVarId, int], flavor: int,
-         upper_bound: Type, *, line: int = -1, column: int = -1
+         upper_bound: Type, *, line: int = -1, column: int = -1, prefix: Optional['Parameters'] = None
     ) -> None:
         super().__init__(name, fullname, id, upper_bound, line=line, column=column)
         self.flavor = flavor
+        self.prefix = prefix or Parameters([], [], [])
 
     @staticmethod
     def new_unification_variable(old: 'ParamSpecType') -> 'ParamSpecType':
         new_id = TypeVarId.new(meta_level=1)
         return ParamSpecType(old.name, old.fullname, new_id, old.flavor, old.upper_bound,
-                             line=old.line, column=old.column)
+                             line=old.line, column=old.column, prefix=old.prefix)
 
     def with_flavor(self, flavor: int) -> 'ParamSpecType':
         return ParamSpecType(self.name, self.fullname, self.id, flavor,
-                             upper_bound=self.upper_bound)
+                             upper_bound=self.upper_bound, prefix=self.prefix)
+
+    def copy_modified(self, *,
+                      id: Bogus[Union[TypeVarId, int]] = _dummy,
+                      flavor: Bogus[int] = _dummy,
+                      prefix: Bogus['Parameters'] = _dummy) -> 'ParamSpecType':
+        return ParamSpecType(
+            self.name,
+            self.fullname,
+            id if id is not _dummy else self.id,
+            flavor if flavor is not _dummy else self.flavor,
+            self.upper_bound,
+            line=self.line,
+            column=self.column,
+            prefix=prefix if prefix is not _dummy else self.prefix,
+        )
 
     def accept(self, visitor: 'TypeVisitor[T]') -> T:
         return visitor.visit_param_spec(self)
@@ -539,6 +556,7 @@ def serialize(self) -> JsonDict:
             'id': self.id.raw_id,
             'flavor': self.flavor,
             'upper_bound': self.upper_bound.serialize(),
+            'prefix': self.prefix.serialize()
         }
 
     @classmethod
@@ -550,6 +568,7 @@ def deserialize(cls, data: JsonDict) -> 'ParamSpecType':
             data['id'],
             data['flavor'],
             deserialize_type(data['upper_bound']),
+            prefix=deserialize_type(data['prefix'])
         )
 
 
@@ -1079,7 +1098,6 @@ def __init__(self,
                  arg_kinds: List[ArgKind],
                  arg_names: Sequence[Optional[str]],
                  ) -> None:
-        #print(f"Parameters.__init__({arg_types=}, {arg_kinds=}, {arg_names=})")
         self.arg_types = list(arg_types)
         self.arg_kinds = arg_kinds
         self.arg_names = list(arg_names)
@@ -1493,7 +1511,7 @@ def param_spec(self) -> Optional[ParamSpecType]:
         if not isinstance(arg_type, ParamSpecType):
             return None
         return ParamSpecType(arg_type.name, arg_type.fullname, arg_type.id, ParamSpecFlavor.BARE,
-                             arg_type.upper_bound)
+                             arg_type.upper_bound, prefix=arg_type.prefix)
 
     def expand_param_spec(self, c: 'CallableType') -> 'CallableType':
         return self.copy_modified(arg_types=self.arg_types[:-2] + c.arg_types,
@@ -2365,12 +2383,18 @@ def visit_type_var(self, t: TypeVarType) -> str:
         return s
 
     def visit_param_spec(self, t: ParamSpecType) -> str:
+        # prefixes are displayed as Concatenate
+        s = ''
+        if t.prefix.arg_types:
+            s += f'Concatenate[{self.list_str(t.prefix.arg_types)}, '
         if t.name is None:
             # Anonymous type variable type (only numeric id).
-            s = f'`{t.id}'
+            s += f'`{t.id}'
         else:
             # Named type variable type.
-            s = f'{t.name_with_suffix()}`{t.id}'
+            s += f'{t.name_with_suffix()}`{t.id}'
+        if t.prefix.arg_types:
+            s += ']'
         return s
 
     def visit_parameters(self, t: Parameters) -> str: