Skip to content

Commit

Permalink
Narrow tuple types using len() (#16237)
Browse files Browse the repository at this point in the history
Fixes #1178 
Supersedes #10367 

This is includes implementation for fixed length tuples, homogeneous
tuples, and variadic tuples (and combinations of those). Generally
implementation is straightforward. Some notes:
* Unfortunately, it is necessary to add a new attribute `min_len` to
`TypeVarTupleType`, which is probably fine, as it doesn't have that many
attributes so far.
* Supporting more general use cases (like `>` comparisons for variadic
tuples) can cause quick proliferation of unions. I added two mechanisms
to counteract this: not applying the narrowing if the integer literal in
comparison is itself large, and collapsing unions of tuples into a
single tuple (if possible) after we are done with the narrowing. This
looks a bit arbitrary, but I think it is important to have.
* Main missing feature here is probably not inferring type information
from indirect comparisons like `len(x) > foo() > 1`. Supporting this
kind of things in full generality is cumbersome, and we may add cases
that turn out to be important later.
* Note I am quite careful with indexing "inside" a `TypeVarTuple`, it is
not really needed now, but I wanted to make everything future proof, so
that it will be easy to add support for upper bounds for
`TypeVarTuple`s, like `Nums = TypeVarTuple("Nums", bound=tuple[float,
...])`.
* I also fix couple existing inconsistencies with `Any` handling in type
narrowing. It looks like they stem from the old incorrect logic that
meet of `Any` and `X` should be `X`, while in fact it should be `Any`.
These fixes are not strictly necessary, but otherwise there may be new
false positives, because I introduce a bunch of additional type
narrowing scenarios here.

cc @hatal175, thanks for the test cases, and for your nice first attempt
to implement this!
Co-authored-by: Tal Hayon <talhayon1@gmail.com>
  • Loading branch information
ilevkivskyi committed Oct 21, 2023
1 parent ff8cebb commit a3af87b
Show file tree
Hide file tree
Showing 16 changed files with 1,154 additions and 38 deletions.
83 changes: 83 additions & 0 deletions mypy/binder.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,17 @@
from mypy.subtypes import is_same_type, is_subtype
from mypy.types import (
AnyType,
Instance,
NoneType,
PartialType,
ProperType,
TupleType,
Type,
TypeOfAny,
TypeType,
UnionType,
UnpackType,
find_unpack_in_list,
get_proper_type,
)
from mypy.typevars import fill_typevars_with_any
Expand Down Expand Up @@ -213,6 +218,24 @@ def update_from_options(self, frames: list[Frame]) -> bool:
for other in resulting_values[1:]:
assert other is not None
type = join_simple(self.declarations[key], type, other)
# Try simplifying resulting type for unions involving variadic tuples.
# Technically, everything is still valid without this step, but if we do
# not do this, this may create long unions after exiting an if check like:
# x: tuple[int, ...]
# if len(x) < 10:
# ...
# We want the type of x to be tuple[int, ...] after this block (if it is
# still equivalent to such type).
if isinstance(type, UnionType):
type = collapse_variadic_union(type)
if isinstance(type, ProperType) and isinstance(type, UnionType):
# Simplify away any extra Any's that were added to the declared
# type when popping a frame.
simplified = UnionType.make_union(
[t for t in type.items if not isinstance(get_proper_type(t), AnyType)]
)
if simplified == self.declarations[key]:
type = simplified
if current_value is None or not is_same_type(type, current_value):
self._put(key, type)
changed = True
Expand Down Expand Up @@ -453,3 +476,63 @@ def get_declaration(expr: BindableExpression) -> Type | None:
elif isinstance(expr.node, TypeInfo):
return TypeType(fill_typevars_with_any(expr.node))
return None


def collapse_variadic_union(typ: UnionType) -> Type:
"""Simplify a union involving variadic tuple if possible.
This will collapse a type like e.g.
tuple[X, Z] | tuple[X, Y, Z] | tuple[X, Y, Y, *tuple[Y, ...], Z]
back to
tuple[X, *tuple[Y, ...], Z]
which is equivalent, but much simpler form of the same type.
"""
tuple_items = []
other_items = []
for t in typ.items:
p_t = get_proper_type(t)
if isinstance(p_t, TupleType):
tuple_items.append(p_t)
else:
other_items.append(t)
if len(tuple_items) <= 1:
# This type cannot be simplified further.
return typ
tuple_items = sorted(tuple_items, key=lambda t: len(t.items))
first = tuple_items[0]
last = tuple_items[-1]
unpack_index = find_unpack_in_list(last.items)
if unpack_index is None:
return typ
unpack = last.items[unpack_index]
assert isinstance(unpack, UnpackType)
unpacked = get_proper_type(unpack.type)
if not isinstance(unpacked, Instance):
return typ
assert unpacked.type.fullname == "builtins.tuple"
suffix = last.items[unpack_index + 1 :]

# Check that first item matches the expected pattern and infer prefix.
if len(first.items) < len(suffix):
return typ
if suffix and first.items[-len(suffix) :] != suffix:
return typ
if suffix:
prefix = first.items[: -len(suffix)]
else:
prefix = first.items

# Check that all middle types match the expected pattern as well.
arg = unpacked.args[0]
for i, it in enumerate(tuple_items[1:-1]):
if it.items != prefix + [arg] * (i + 1) + suffix:
return typ

# Check the last item (the one with unpack), and choose an appropriate simplified type.
if last.items != prefix + [arg] * (len(typ.items) - 1) + [unpack] + suffix:
return typ
if len(first.items) == 0:
simplified: Type = unpacked.copy_modified()
else:
simplified = TupleType(prefix + [unpack] + suffix, fallback=last.partial_fallback)
return UnionType.make_union([simplified] + other_items)
Loading

0 comments on commit a3af87b

Please sign in to comment.