|
12 | 12 | from mypy.subtypes import is_same_type, is_subtype |
13 | 13 | from mypy.types import ( |
14 | 14 | AnyType, |
| 15 | + Instance, |
15 | 16 | NoneType, |
16 | 17 | PartialType, |
| 18 | + ProperType, |
| 19 | + TupleType, |
17 | 20 | Type, |
18 | 21 | TypeOfAny, |
19 | 22 | TypeType, |
20 | 23 | UnionType, |
| 24 | + UnpackType, |
| 25 | + find_unpack_in_list, |
21 | 26 | get_proper_type, |
22 | 27 | ) |
23 | 28 | from mypy.typevars import fill_typevars_with_any |
@@ -213,6 +218,24 @@ def update_from_options(self, frames: list[Frame]) -> bool: |
213 | 218 | for other in resulting_values[1:]: |
214 | 219 | assert other is not None |
215 | 220 | type = join_simple(self.declarations[key], type, other) |
| 221 | + # Try simplifying resulting type for unions involving variadic tuples. |
| 222 | + # Technically, everything is still valid without this step, but if we do |
| 223 | + # not do this, this may create long unions after exiting an if check like: |
| 224 | + # x: tuple[int, ...] |
| 225 | + # if len(x) < 10: |
| 226 | + # ... |
| 227 | + # We want the type of x to be tuple[int, ...] after this block (if it is |
| 228 | + # still equivalent to such type). |
| 229 | + if isinstance(type, UnionType): |
| 230 | + type = collapse_variadic_union(type) |
| 231 | + if isinstance(type, ProperType) and isinstance(type, UnionType): |
| 232 | + # Simplify away any extra Any's that were added to the declared |
| 233 | + # type when popping a frame. |
| 234 | + simplified = UnionType.make_union( |
| 235 | + [t for t in type.items if not isinstance(get_proper_type(t), AnyType)] |
| 236 | + ) |
| 237 | + if simplified == self.declarations[key]: |
| 238 | + type = simplified |
216 | 239 | if current_value is None or not is_same_type(type, current_value): |
217 | 240 | self._put(key, type) |
218 | 241 | changed = True |
@@ -453,3 +476,63 @@ def get_declaration(expr: BindableExpression) -> Type | None: |
453 | 476 | elif isinstance(expr.node, TypeInfo): |
454 | 477 | return TypeType(fill_typevars_with_any(expr.node)) |
455 | 478 | return None |
| 479 | + |
| 480 | + |
| 481 | +def collapse_variadic_union(typ: UnionType) -> Type: |
| 482 | + """Simplify a union involving variadic tuple if possible. |
| 483 | +
|
| 484 | + This will collapse a type like e.g. |
| 485 | + tuple[X, Z] | tuple[X, Y, Z] | tuple[X, Y, Y, *tuple[Y, ...], Z] |
| 486 | + back to |
| 487 | + tuple[X, *tuple[Y, ...], Z] |
| 488 | + which is equivalent, but much simpler form of the same type. |
| 489 | + """ |
| 490 | + tuple_items = [] |
| 491 | + other_items = [] |
| 492 | + for t in typ.items: |
| 493 | + p_t = get_proper_type(t) |
| 494 | + if isinstance(p_t, TupleType): |
| 495 | + tuple_items.append(p_t) |
| 496 | + else: |
| 497 | + other_items.append(t) |
| 498 | + if len(tuple_items) <= 1: |
| 499 | + # This type cannot be simplified further. |
| 500 | + return typ |
| 501 | + tuple_items = sorted(tuple_items, key=lambda t: len(t.items)) |
| 502 | + first = tuple_items[0] |
| 503 | + last = tuple_items[-1] |
| 504 | + unpack_index = find_unpack_in_list(last.items) |
| 505 | + if unpack_index is None: |
| 506 | + return typ |
| 507 | + unpack = last.items[unpack_index] |
| 508 | + assert isinstance(unpack, UnpackType) |
| 509 | + unpacked = get_proper_type(unpack.type) |
| 510 | + if not isinstance(unpacked, Instance): |
| 511 | + return typ |
| 512 | + assert unpacked.type.fullname == "builtins.tuple" |
| 513 | + suffix = last.items[unpack_index + 1 :] |
| 514 | + |
| 515 | + # Check that first item matches the expected pattern and infer prefix. |
| 516 | + if len(first.items) < len(suffix): |
| 517 | + return typ |
| 518 | + if suffix and first.items[-len(suffix) :] != suffix: |
| 519 | + return typ |
| 520 | + if suffix: |
| 521 | + prefix = first.items[: -len(suffix)] |
| 522 | + else: |
| 523 | + prefix = first.items |
| 524 | + |
| 525 | + # Check that all middle types match the expected pattern as well. |
| 526 | + arg = unpacked.args[0] |
| 527 | + for i, it in enumerate(tuple_items[1:-1]): |
| 528 | + if it.items != prefix + [arg] * (i + 1) + suffix: |
| 529 | + return typ |
| 530 | + |
| 531 | + # Check the last item (the one with unpack), and choose an appropriate simplified type. |
| 532 | + if last.items != prefix + [arg] * (len(typ.items) - 1) + [unpack] + suffix: |
| 533 | + return typ |
| 534 | + if len(first.items) == 0: |
| 535 | + simplified: Type = unpacked.copy_modified() |
| 536 | + else: |
| 537 | + simplified = TupleType(prefix + [unpack] + suffix, fallback=last.partial_fallback) |
| 538 | + return UnionType.make_union([simplified] + other_items) |
0 commit comments