Skip to content

Commit

Permalink
Support user defined variadic tuple types (#15961)
Browse files Browse the repository at this point in the history
Fixes #15946

Note this actually adds support also for variadic NamedTuples and
variadic TypedDicts. Not that anyone requested this, but since generic
NamedTuples and generic TypedDicts are supported using the same
mechanism (special aliases) as generic tuple types (like `class
A(Tuple[T, S]): ...` in the issue), it looked more risky and arbitrary
to _not_support them.

Btw the implementation is simple, but while I was working on this, I
accidentally found a problem with my general idea of doing certain type
normlaizations in `semanal_typeargs.py`. The problem is that sometimes
we can call `get_proper_type()` during semantic analysis, so all the
code that gets triggered by this (mostly `expand_type()`) can't really
rely on types being normalized. Fortunately, with just few tweaks I
manged to make the code mostly robust to such scenarios (TBH there are
few possible holes left, but this is getting really complex, I think it
is better to release this, and see if people will ever hit such
scenarios, then fix accordingly).
  • Loading branch information
ilevkivskyi authored Aug 26, 2023
1 parent 29abf39 commit efecd59
Show file tree
Hide file tree
Showing 9 changed files with 149 additions and 21 deletions.
7 changes: 6 additions & 1 deletion mypy/expandtype.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,8 @@ def visit_unpack_type(self, t: UnpackType) -> Type:
# instead.
# However, if the item is a variadic tuple, we can simply carry it over.
# In particular, if we expand A[*tuple[T, ...]] with substitutions {T: str},
# it is hard to assert this without getting proper type.
# it is hard to assert this without getting proper type. Another important
# example is non-normalized types when called from semanal.py.
return UnpackType(t.type.accept(self))

def expand_unpack(self, t: UnpackType) -> list[Type] | AnyType | UninhabitedType:
Expand Down Expand Up @@ -414,6 +415,10 @@ def visit_tuple_type(self, t: TupleType) -> Type:
unpacked = get_proper_type(item.type)
if isinstance(unpacked, Instance):
assert unpacked.type.fullname == "builtins.tuple"
if t.partial_fallback.type.fullname != "builtins.tuple":
# If it is a subtype (like named tuple) we need to preserve it,
# this essentially mimics the logic in tuple_fallback().
return t.partial_fallback.accept(self)
return unpacked
fallback = t.partial_fallback.accept(self)
assert isinstance(fallback, ProperType) and isinstance(fallback, Instance)
Expand Down
1 change: 0 additions & 1 deletion mypy/maptype.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,5 @@ def instance_to_type_environment(instance: Instance) -> dict[TypeVarId, Type]:
required number of type arguments. So this environment consists
of the class's type variables mapped to the Instance's actual
arguments. The type variables are mapped by their `id`.
"""
return {binder.id: arg for binder, arg in zip(instance.type.defn.type_vars, instance.args)}
12 changes: 10 additions & 2 deletions mypy/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -3546,7 +3546,12 @@ def from_tuple_type(cls, info: TypeInfo) -> TypeAlias:
assert info.tuple_type
# TODO: is it possible to refactor this to set the correct type vars here?
return TypeAlias(
info.tuple_type.copy_modified(fallback=mypy.types.Instance(info, info.defn.type_vars)),
info.tuple_type.copy_modified(
# Create an Instance similar to fill_typevars().
fallback=mypy.types.Instance(
info, mypy.types.type_vars_as_args(info.defn.type_vars)
)
),
info.fullname,
info.line,
info.column,
Expand All @@ -3563,7 +3568,10 @@ def from_typeddict_type(cls, info: TypeInfo) -> TypeAlias:
# TODO: is it possible to refactor this to set the correct type vars here?
return TypeAlias(
info.typeddict_type.copy_modified(
fallback=mypy.types.Instance(info, info.defn.type_vars)
# Create an Instance similar to fill_typevars().
fallback=mypy.types.Instance(
info, mypy.types.type_vars_as_args(info.defn.type_vars)
)
),
info.fullname,
info.line,
Expand Down
10 changes: 8 additions & 2 deletions mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@
get_proper_types,
is_named_instance,
remove_dups,
type_vars_as_args,
)
from mypy.types_utils import is_invalid_recursive_alias, store_argument_type
from mypy.typevars import fill_typevars
Expand Down Expand Up @@ -1702,12 +1703,17 @@ def setup_type_vars(self, defn: ClassDef, tvar_defs: list[TypeVarLikeType]) -> N
def setup_alias_type_vars(self, defn: ClassDef) -> None:
assert defn.info.special_alias is not None
defn.info.special_alias.alias_tvars = list(defn.type_vars)
# It is a bit unfortunate that we need to inline some logic from TypeAlias constructor,
# but it is required, since type variables may change during semantic analyzer passes.
for i, t in enumerate(defn.type_vars):
if isinstance(t, TypeVarTupleType):
defn.info.special_alias.tvar_tuple_index = i
target = defn.info.special_alias.target
assert isinstance(target, ProperType)
if isinstance(target, TypedDictType):
target.fallback.args = tuple(defn.type_vars)
target.fallback.args = type_vars_as_args(defn.type_vars)
elif isinstance(target, TupleType):
target.partial_fallback.args = tuple(defn.type_vars)
target.partial_fallback.args = type_vars_as_args(defn.type_vars)
else:
assert False, f"Unexpected special alias type: {type(target)}"

Expand Down
14 changes: 7 additions & 7 deletions mypy/semanal_typeargs.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,31 +86,31 @@ def visit_type_alias_type(self, t: TypeAliasType) -> None:
# correct aliases. Also, variadic aliases are better to check when fully analyzed,
# so we do this here.
assert t.alias is not None, f"Unfixed type alias {t.type_ref}"
args = flatten_nested_tuples(t.args)
# TODO: consider moving this validation to typeanal.py, expanding invalid aliases
# during semantic analysis may cause crashes.
if t.alias.tvar_tuple_index is not None:
correct = len(args) >= len(t.alias.alias_tvars) - 1
correct = len(t.args) >= len(t.alias.alias_tvars) - 1
if any(
isinstance(a, UnpackType) and isinstance(get_proper_type(a.type), Instance)
for a in args
for a in t.args
):
correct = True
else:
correct = len(args) == len(t.alias.alias_tvars)
correct = len(t.args) == len(t.alias.alias_tvars)
if not correct:
if t.alias.tvar_tuple_index is not None:
exp_len = f"at least {len(t.alias.alias_tvars) - 1}"
else:
exp_len = f"{len(t.alias.alias_tvars)}"
self.fail(
f"Bad number of arguments for type alias, expected: {exp_len}, given: {len(args)}",
"Bad number of arguments for type alias,"
f" expected: {exp_len}, given: {len(t.args)}",
t,
code=codes.TYPE_ARG,
)
t.args = set_any_tvars(
t.alias, t.line, t.column, self.options, from_error=True, fail=self.fail
).args
else:
t.args = args
is_error = self.validate_args(t.alias.name, t.args, t.alias.alias_tvars, t)
if not is_error:
# If there was already an error for the alias itself, there is no point in checking
Expand Down
1 change: 1 addition & 0 deletions mypy/semanal_typeddict.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,7 @@ def map_items_to_base(
if not tvars:
mapped_items[key] = type_in_base
continue
# TODO: simple zip can't be used for variadic types.
mapped_items[key] = expand_type(
type_in_base, {t.id: a for (t, a) in zip(tvars, base_args)}
)
Expand Down
12 changes: 8 additions & 4 deletions mypy/typeanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@
UnionType,
UnpackType,
callable_with_ellipsis,
flatten_nested_tuples,
flatten_nested_unions,
get_proper_type,
has_type_vars,
Expand Down Expand Up @@ -763,8 +764,8 @@ def analyze_type_with_type_info(
if info.special_alias:
return instantiate_type_alias(
info.special_alias,
# TODO: should we allow NamedTuples generic in ParamSpec and TypeVarTuple?
self.anal_array(args),
# TODO: should we allow NamedTuples generic in ParamSpec?
self.anal_array(args, allow_unpack=True),
self.fail,
False,
ctx,
Expand All @@ -782,7 +783,7 @@ def analyze_type_with_type_info(
return instantiate_type_alias(
info.special_alias,
# TODO: should we allow TypedDicts generic in ParamSpec?
self.anal_array(args),
self.anal_array(args, allow_unpack=True),
self.fail,
False,
ctx,
Expand Down Expand Up @@ -1948,7 +1949,10 @@ def instantiate_type_alias(
# TODO: we need to check args validity w.r.t alias.alias_tvars.
# Otherwise invalid instantiations will be allowed in runtime context.
# Note: in type context, these will be still caught by semanal_typeargs.
typ = TypeAliasType(node, args, ctx.line, ctx.column)
# Type aliases are special, since they can be expanded during semantic analysis,
# so we need to normalize them as soon as possible.
# TODO: can this cause an infinite recursion?
typ = TypeAliasType(node, flatten_nested_tuples(args), ctx.line, ctx.column)
assert typ.alias is not None
# HACK: Implement FlexibleAlias[T, typ] by expanding it to typ here.
if (
Expand Down
26 changes: 22 additions & 4 deletions mypy/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -1045,9 +1045,12 @@ class UnpackType(ProperType):
"""Type operator Unpack from PEP646. Can be either with Unpack[]
or unpacking * syntax.
The inner type should be either a TypeVarTuple, a constant size
tuple, or a variable length tuple. Type aliases to these are not allowed,
except during semantic analysis.
The inner type should be either a TypeVarTuple, or a variable length tuple.
In an exceptional case of callable star argument it can be a fixed length tuple.
Note: the above restrictions are only guaranteed by normalizations after semantic
analysis, if your code needs to handle UnpackType *during* semantic analysis, it is
wild west, technically anything can be present in the wrapped type.
"""

__slots__ = ["type"]
Expand Down Expand Up @@ -2143,7 +2146,11 @@ def with_normalized_var_args(self) -> Self:
assert nested_unpacked.type.fullname == "builtins.tuple"
new_unpack = nested_unpacked.args[0]
else:
assert isinstance(nested_unpacked, TypeVarTupleType)
if not isinstance(nested_unpacked, TypeVarTupleType):
# We found a non-nomralized tuple type, this means this method
# is called during semantic analysis (e.g. from get_proper_type())
# there is no point in normalizing callables at this stage.
return self
new_unpack = nested_unpack
else:
new_unpack = UnpackType(
Expand Down Expand Up @@ -3587,6 +3594,17 @@ def remove_dups(types: list[T]) -> list[T]:
return new_types


def type_vars_as_args(type_vars: Sequence[TypeVarLikeType]) -> tuple[Type, ...]:
"""Represent type variables as they would appear in a type argument list."""
args: list[Type] = []
for tv in type_vars:
if isinstance(tv, TypeVarTupleType):
args.append(UnpackType(tv))
else:
args.append(tv)
return tuple(args)


# This cyclic import is unfortunate, but to avoid it we would need to move away all uses
# of get_proper_type() from types.py. Majority of them have been removed, but few remaining
# are quite tricky to get rid of, but ultimately we want to do it at some point.
Expand Down
87 changes: 87 additions & 0 deletions test-data/unit/check-typevar-tuple.test
Original file line number Diff line number Diff line change
Expand Up @@ -1032,3 +1032,90 @@ Second = Tuple[C, D]
x: G[Unpack[First], Unpack[Second]] # E: Type argument "A" of "G" must be a subtype of "int" \
# E: Type argument "D" of "G" must be a subtype of "str"
[builtins fixtures/tuple.pyi]

[case testVariadicTupleType]
from typing import Tuple, Callable
from typing_extensions import TypeVarTuple, Unpack

Ts = TypeVarTuple("Ts")
class A(Tuple[Unpack[Ts]]):
fn: Callable[[Unpack[Ts]], None]

x: A[int]
reveal_type(x) # N: Revealed type is "Tuple[builtins.int, fallback=__main__.A[builtins.int]]"
reveal_type(x[0]) # N: Revealed type is "builtins.int"
reveal_type(x.fn) # N: Revealed type is "def (builtins.int)"

y: A[int, str]
reveal_type(y) # N: Revealed type is "Tuple[builtins.int, builtins.str, fallback=__main__.A[builtins.int, builtins.str]]"
reveal_type(y[0]) # N: Revealed type is "builtins.int"
reveal_type(y.fn) # N: Revealed type is "def (builtins.int, builtins.str)"

z: A[Unpack[Tuple[int, ...]]]
reveal_type(z) # N: Revealed type is "__main__.A[Unpack[builtins.tuple[builtins.int, ...]]]"
# TODO: this requires fixing map_instance_to_supertype().
# reveal_type(z[0])
reveal_type(z.fn) # N: Revealed type is "def (*builtins.int)"

t: A[int, Unpack[Tuple[int, str]], str]
reveal_type(t) # N: Revealed type is "Tuple[builtins.int, builtins.int, builtins.str, builtins.str, fallback=__main__.A[builtins.int, builtins.int, builtins.str, builtins.str]]"
reveal_type(t[0]) # N: Revealed type is "builtins.int"
reveal_type(t.fn) # N: Revealed type is "def (builtins.int, builtins.int, builtins.str, builtins.str)"
[builtins fixtures/tuple.pyi]

[case testVariadicNamedTuple]
from typing import Tuple, Callable, NamedTuple, Generic
from typing_extensions import TypeVarTuple, Unpack

Ts = TypeVarTuple("Ts")
class A(NamedTuple, Generic[Unpack[Ts], T]):
fn: Callable[[Unpack[Ts]], None]
val: T

y: A[int, str]
reveal_type(y) # N: Revealed type is "Tuple[def (builtins.int), builtins.str, fallback=__main__.A[builtins.int, builtins.str]]"
reveal_type(y[0]) # N: Revealed type is "def (builtins.int)"
reveal_type(y.fn) # N: Revealed type is "def (builtins.int)"

z: A[Unpack[Tuple[int, ...]]]
reveal_type(z) # N: Revealed type is "Tuple[def (*builtins.int), builtins.int, fallback=__main__.A[Unpack[builtins.tuple[builtins.int, ...]], builtins.int]]"
reveal_type(z.fn) # N: Revealed type is "def (*builtins.int)"

t: A[int, Unpack[Tuple[int, str]], str]
reveal_type(t) # N: Revealed type is "Tuple[def (builtins.int, builtins.int, builtins.str), builtins.str, fallback=__main__.A[builtins.int, builtins.int, builtins.str, builtins.str]]"

def test(x: int, y: str) -> None: ...
nt = A(fn=test, val=42)
reveal_type(nt) # N: Revealed type is "Tuple[def (builtins.int, builtins.str), builtins.int, fallback=__main__.A[builtins.int, builtins.str, builtins.int]]"

def bad() -> int: ...
nt2 = A(fn=bad, val=42) # E: Argument "fn" to "A" has incompatible type "Callable[[], int]"; expected "Callable[[], None]"
[builtins fixtures/tuple.pyi]

[case testVariadicTypedDict]
from typing import Tuple, Callable, Generic
from typing_extensions import TypeVarTuple, Unpack, TypedDict

Ts = TypeVarTuple("Ts")
class A(TypedDict, Generic[Unpack[Ts], T]):
fn: Callable[[Unpack[Ts]], None]
val: T

y: A[int, str]
reveal_type(y) # N: Revealed type is "TypedDict('__main__.A', {'fn': def (builtins.int), 'val': builtins.str})"
reveal_type(y["fn"]) # N: Revealed type is "def (builtins.int)"

z: A[Unpack[Tuple[int, ...]]]
reveal_type(z) # N: Revealed type is "TypedDict('__main__.A', {'fn': def (*builtins.int), 'val': builtins.int})"
reveal_type(z["fn"]) # N: Revealed type is "def (*builtins.int)"

t: A[int, Unpack[Tuple[int, str]], str]
reveal_type(t) # N: Revealed type is "TypedDict('__main__.A', {'fn': def (builtins.int, builtins.int, builtins.str), 'val': builtins.str})"

def test(x: int, y: str) -> None: ...
td = A({"fn": test, "val": 42})
reveal_type(td) # N: Revealed type is "TypedDict('__main__.A', {'fn': def (builtins.int, builtins.str), 'val': builtins.int})"

def bad() -> int: ...
td2 = A({"fn": bad, "val": 42}) # E: Incompatible types (expression has type "Callable[[], int]", TypedDict item "fn" has type "Callable[[], None]")
[builtins fixtures/tuple.pyi]

0 comments on commit efecd59

Please sign in to comment.