Support PEP-646 and PEP-692 in the same callable (#16294)

ilevkivskyi · web-flow · commit 5d4046477eb0 · 2023-10-27T18:36:08.000+01:00
Fixes #16285 I was not sure if it is important to support this, but taking into account the current behavior is a crash, and that implementation is quite simple, I think we should do this. Using this opportunity I also improve related error messages a bit.
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -950,7 +950,7 @@ def remove_unpack_kwargs(self, defn: FuncDef, typ: CallableType) -> CallableType
             return typ
         last_type = get_proper_type(last_type.type)
         if not isinstance(last_type, TypedDictType):
-            self.fail("Unpack item in ** argument must be a TypedDict", defn)
+            self.fail("Unpack item in ** argument must be a TypedDict", last_type)
             new_arg_types = typ.arg_types[:-1] + [AnyType(TypeOfAny.from_error)]
             return typ.copy_modified(arg_types=new_arg_types)
         overlap = set(typ.arg_names) & set(last_type.items)
diff --git a/mypy/typeanal.py b/mypy/typeanal.py
@@ -987,33 +987,40 @@ def visit_callable_type(self, t: CallableType, nested: bool = True) -> Type:
                     self.anal_star_arg_type(t.arg_types[-2], ARG_STAR, nested=nested),
                     self.anal_star_arg_type(t.arg_types[-1], ARG_STAR2, nested=nested),
                 ]
+                # If nested is True, it means we are analyzing a Callable[...] type, rather
+                # than a function definition type. We need to "unpack" ** TypedDict annotation
+                # here (for function definitions it is done in semanal).
+                if nested and isinstance(arg_types[-1], UnpackType):
+                    # TODO: it would be better to avoid this get_proper_type() call.
+                    unpacked = get_proper_type(arg_types[-1].type)
+                    if isinstance(unpacked, TypedDictType):
+                        arg_types[-1] = unpacked
+                        unpacked_kwargs = True
+                    arg_types = self.check_unpacks_in_list(arg_types)
             else:
-                arg_types = self.anal_array(t.arg_types, nested=nested, allow_unpack=True)
                 star_index = None
                 if ARG_STAR in arg_kinds:
                     star_index = arg_kinds.index(ARG_STAR)
                 star2_index = None
                 if ARG_STAR2 in arg_kinds:
                     star2_index = arg_kinds.index(ARG_STAR2)
-                validated_args: list[Type] = []
-                for i, at in enumerate(arg_types):
-                    if isinstance(at, UnpackType) and i not in (star_index, star2_index):
-                        self.fail(
-                            message_registry.INVALID_UNPACK_POSITION, at, code=codes.VALID_TYPE
-                        )
-                        validated_args.append(AnyType(TypeOfAny.from_error))
-                    else:
-                        if nested and isinstance(at, UnpackType) and i == star_index:
-                            # TODO: it would be better to avoid this get_proper_type() call.
-                            p_at = get_proper_type(at.type)
-                            if isinstance(p_at, TypedDictType) and not at.from_star_syntax:
-                                # Automatically detect Unpack[Foo] in Callable as backwards
-                                # compatible syntax for **Foo, if Foo is a TypedDict.
-                                at = p_at
-                                arg_kinds[i] = ARG_STAR2
-                                unpacked_kwargs = True
-                        validated_args.append(at)
-                arg_types = validated_args
+                arg_types = []
+                for i, ut in enumerate(t.arg_types):
+                    at = self.anal_type(
+                        ut, nested=nested, allow_unpack=i in (star_index, star2_index)
+                    )
+                    if nested and isinstance(at, UnpackType) and i == star_index:
+                        # TODO: it would be better to avoid this get_proper_type() call.
+                        p_at = get_proper_type(at.type)
+                        if isinstance(p_at, TypedDictType) and not at.from_star_syntax:
+                            # Automatically detect Unpack[Foo] in Callable as backwards
+                            # compatible syntax for **Foo, if Foo is a TypedDict.
+                            at = p_at
+                            arg_kinds[i] = ARG_STAR2
+                            unpacked_kwargs = True
+                    arg_types.append(at)
+                if nested:
+                    arg_types = self.check_unpacks_in_list(arg_types)
             # If there were multiple (invalid) unpacks, the arg types list will become shorter,
             # we need to trim the kinds/names as well to avoid crashes.
             arg_kinds = t.arg_kinds[: len(arg_types)]
@@ -1387,8 +1394,9 @@ def analyze_callable_args(
         names: list[str | None] = []
         seen_unpack = False
         unpack_types: list[Type] = []
-        invalid_unpacks = []
-        for arg in arglist.items:
+        invalid_unpacks: list[Type] = []
+        second_unpack_last = False
+        for i, arg in enumerate(arglist.items):
             if isinstance(arg, CallableArgument):
                 args.append(arg.typ)
                 names.append(arg.name)
@@ -1415,6 +1423,11 @@ def analyze_callable_args(
             ):
                 if seen_unpack:
                     # Multiple unpacks, preserve them, so we can give an error later.
+                    if i == len(arglist.items) - 1 and not invalid_unpacks:
+                        # Special case: if there are just two unpacks, and the second one appears
+                        # as last type argument, it can be still valid, if the second unpacked type
+                        # is a TypedDict. This should be checked by the caller.
+                        second_unpack_last = True
                     invalid_unpacks.append(arg)
                     continue
                 seen_unpack = True
@@ -1442,7 +1455,7 @@ def analyze_callable_args(
             names.append(None)
         for arg in invalid_unpacks:
             args.append(arg)
-            kinds.append(ARG_STAR)
+            kinds.append(ARG_STAR2 if second_unpack_last else ARG_STAR)
             names.append(None)
         # Note that arglist below is only used for error context.
         check_arg_names(names, [arglist] * len(args), self.fail, "Callable")
diff --git a/mypy/types.py b/mypy/types.py
@@ -3268,15 +3268,16 @@ def visit_callable_type(self, t: CallableType) -> str:
             num_skip = 0
 
         s = ""
-        bare_asterisk = False
+        asterisk = False
         for i in range(len(t.arg_types) - num_skip):
             if s != "":
                 s += ", "
-            if t.arg_kinds[i].is_named() and not bare_asterisk:
+            if t.arg_kinds[i].is_named() and not asterisk:
                 s += "*, "
-                bare_asterisk = True
+                asterisk = True
             if t.arg_kinds[i] == ARG_STAR:
                 s += "*"
+                asterisk = True
             if t.arg_kinds[i] == ARG_STAR2:
                 s += "**"
             name = t.arg_names[i]
diff --git a/test-data/unit/check-typevar-tuple.test b/test-data/unit/check-typevar-tuple.test
@@ -571,8 +571,7 @@ from typing_extensions import Unpack, TypeVarTuple
 
 Ts = TypeVarTuple("Ts")
 Us = TypeVarTuple("Us")
-a: Callable[[Unpack[Ts], Unpack[Us]], int]  # E: Var args may not appear after named or var args \
-                                            # E: More than one Unpack in a type is not allowed
+a: Callable[[Unpack[Ts], Unpack[Us]], int]  # E: More than one Unpack in a type is not allowed
 reveal_type(a)  # N: Revealed type is "def [Ts, Us] (*Unpack[Ts`-1]) -> builtins.int"
 b: Callable[[Unpack], int]  # E: Unpack[...] requires exactly one type argument
 reveal_type(b)  # N: Revealed type is "def (*Any) -> builtins.int"
@@ -730,8 +729,7 @@ A = Tuple[Unpack[Ts], Unpack[Us]]  # E: More than one Unpack in a type is not al
 x: A[int, str]
 reveal_type(x)  # N: Revealed type is "Tuple[builtins.int, builtins.str]"
 
-B = Callable[[Unpack[Ts], Unpack[Us]], int]  # E: Var args may not appear after named or var args \
-                                             # E: More than one Unpack in a type is not allowed
+B = Callable[[Unpack[Ts], Unpack[Us]], int]  # E: More than one Unpack in a type is not allowed
 y: B[int, str]
 reveal_type(y)  # N: Revealed type is "def (builtins.int, builtins.str) -> builtins.int"
 
@@ -1912,3 +1910,101 @@ reveal_type(y)  # N: Revealed type is "__main__.C[builtins.int, Unpack[builtins.
 z = C[int]()  # E: Bad number of arguments, expected: at least 2, given: 1
 reveal_type(z)  # N: Revealed type is "__main__.C[Any, Unpack[builtins.tuple[Any, ...]], Any]"
 [builtins fixtures/tuple.pyi]
+
+[case testTypeVarTupleBothUnpacksSimple]
+from typing import Tuple
+from typing_extensions import Unpack, TypeVarTuple, TypedDict
+
+class Keywords(TypedDict):
+    a: str
+    b: str
+
+Ints = Tuple[int, ...]
+
+def f(*args: Unpack[Ints], other: str = "no", **kwargs: Unpack[Keywords]) -> None: ...
+reveal_type(f)  # N: Revealed type is "def (*args: builtins.int, other: builtins.str =, **kwargs: Unpack[TypedDict('__main__.Keywords', {'a': builtins.str, 'b': builtins.str})])"
+f(1, 2, a="a", b="b")  # OK
+f(1, 2, 3)  # E: Missing named argument "a" for "f" \
+            # E: Missing named argument "b" for "f"
+
+Ts = TypeVarTuple("Ts")
+def g(*args: Unpack[Ts], other: str = "no", **kwargs: Unpack[Keywords]) -> None: ...
+reveal_type(g)  # N: Revealed type is "def [Ts] (*args: Unpack[Ts`-1], other: builtins.str =, **kwargs: Unpack[TypedDict('__main__.Keywords', {'a': builtins.str, 'b': builtins.str})])"
+g(1, 2, a="a", b="b")  # OK
+g(1, 2, 3)  # E: Missing named argument "a" for "g" \
+            # E: Missing named argument "b" for "g"
+
+def bad(
+    *args: Unpack[Keywords],  # E: "Keywords" cannot be unpacked (must be tuple or TypeVarTuple)
+    **kwargs: Unpack[Ints],  # E: Unpack item in ** argument must be a TypedDict
+) -> None: ...
+reveal_type(bad)  # N: Revealed type is "def (*args: Any, **kwargs: Any)"
+
+def bad2(
+    one: int,
+    *args: Unpack[Keywords],  # E: "Keywords" cannot be unpacked (must be tuple or TypeVarTuple)
+    other: str = "no",
+     **kwargs: Unpack[Ints],  # E: Unpack item in ** argument must be a TypedDict
+) -> None: ...
+reveal_type(bad2)  # N: Revealed type is "def (one: builtins.int, *args: Any, other: builtins.str =, **kwargs: Any)"
+[builtins fixtures/tuple.pyi]
+
+[case testTypeVarTupleBothUnpacksCallable]
+from typing import Callable, Tuple
+from typing_extensions import Unpack, TypedDict
+
+class Keywords(TypedDict):
+    a: str
+    b: str
+Ints = Tuple[int, ...]
+
+cb: Callable[[Unpack[Ints], Unpack[Keywords]], None]
+reveal_type(cb)  # N: Revealed type is "def (*builtins.int, **Unpack[TypedDict('__main__.Keywords', {'a': builtins.str, 'b': builtins.str})])"
+
+cb2: Callable[[int, Unpack[Ints], int, Unpack[Keywords]], None]
+reveal_type(cb2)  # N: Revealed type is "def (builtins.int, *Unpack[Tuple[Unpack[builtins.tuple[builtins.int, ...]], builtins.int]], **Unpack[TypedDict('__main__.Keywords', {'a': builtins.str, 'b': builtins.str})])"
+cb2(1, 2, 3, a="a", b="b")
+cb2(1, a="a", b="b")  # E: Too few arguments
+cb2(1, 2, 3, a="a")  # E: Missing named argument "b"
+
+bad1: Callable[[Unpack[Ints], Unpack[Ints]], None]  # E: More than one Unpack in a type is not allowed
+reveal_type(bad1)  # N: Revealed type is "def (*builtins.int)"
+bad2: Callable[[Unpack[Keywords], Unpack[Keywords]], None]  # E: "Keywords" cannot be unpacked (must be tuple or TypeVarTuple)
+reveal_type(bad2)  # N: Revealed type is "def (*Any, **Unpack[TypedDict('__main__.Keywords', {'a': builtins.str, 'b': builtins.str})])"
+bad3: Callable[[Unpack[Keywords], Unpack[Ints]], None]  # E: "Keywords" cannot be unpacked (must be tuple or TypeVarTuple) \
+                                                        # E: More than one Unpack in a type is not allowed
+reveal_type(bad3)  # N: Revealed type is "def (*Any)"
+[builtins fixtures/tuple.pyi]
+
+[case testTypeVarTupleBothUnpacksApplication]
+from typing import Callable, TypeVar, Optional
+from typing_extensions import Unpack, TypeVarTuple, TypedDict
+
+class Keywords(TypedDict):
+    a: str
+    b: str
+
+T = TypeVar("T")
+Ts = TypeVarTuple("Ts")
+def test(
+    x: int,
+    func: Callable[[Unpack[Ts]], T],
+    *args: Unpack[Ts],
+    other: Optional[str] = None,
+    **kwargs: Unpack[Keywords],
+) -> T:
+    if bool():
+        func(*args, **kwargs)  # E: Extra argument "a" from **args
+    return func(*args)
+def test2(
+    x: int,
+    func: Callable[[Unpack[Ts], Unpack[Keywords]], T],
+    *args: Unpack[Ts],
+    other: Optional[str] = None,
+    **kwargs: Unpack[Keywords],
+) -> T:
+    if bool():
+        func(*args)  # E: Missing named argument "a" \
+                     # E: Missing named argument "b"
+    return func(*args, **kwargs)
+[builtins fixtures/tuple.pyi]
diff --git a/test-data/unit/semanal-types.test b/test-data/unit/semanal-types.test
@@ -1043,7 +1043,7 @@ MypyFile:1(
       default(
         Var(y)
         StrExpr()))
-    def (*x: builtins.int, *, y: builtins.str =) -> Any
+    def (*x: builtins.int, y: builtins.str =) -> Any
     VarArg(
       Var(x))
     Block:1(
