From ef8281b6957548112e2d1b3189833b63ea5682a4 Mon Sep 17 00:00:00 2001 From: Suneet Tipirneni <77477100+suneettipirneni@users.noreply.github.com> Date: Mon, 23 Jun 2025 12:38:01 -0400 Subject: [PATCH] [ty] add support for mapped union and intersection subscript loads (#18846) ## Summary Note this modifies the diagnostics a bit. Previously performing subscript access on something like `NotSubscriptable1 | NotSubscriptable2` would report the full type as not being subscriptable: ``` [non-subscriptable] "Cannot subscript object of type `NotSubscriptable1 | NotSubscriptable2` with no `__getitem__` method" ``` Now each erroneous constituent has a separate error: ``` [non-subscriptable] "Cannot subscript object of type `NotSubscriptable2` with no `__getitem__` method" [non-subscriptable] "Cannot subscript object of type `NotSubscriptable1` with no `__getitem__` method" ``` Closes https://github.com/astral-sh/ty/issues/625 ## Test Plan mdtest --------- Co-authored-by: Carl Meyer --- .../resources/mdtest/call/dunder.md | 3 +- .../resources/mdtest/narrow/type_guards.md | 2 +- .../resources/mdtest/subscript/class.md | 11 +++--- .../resources/mdtest/subscript/tuple.md | 27 ++++++++++++++ crates/ty_python_semantic/src/types/infer.rs | 35 ++++++++++--------- 5 files changed, 54 insertions(+), 24 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/call/dunder.md b/crates/ty_python_semantic/resources/mdtest/call/dunder.md index 15cf01b7a1..ed846ea8df 100644 --- a/crates/ty_python_semantic/resources/mdtest/call/dunder.md +++ b/crates/ty_python_semantic/resources/mdtest/call/dunder.md @@ -258,7 +258,8 @@ class NotSubscriptable2: self.__getitem__ = external_getitem def _(union: NotSubscriptable1 | NotSubscriptable2): - # error: [non-subscriptable] + # error: [non-subscriptable] "Cannot subscript object of type `NotSubscriptable2` with no `__getitem__` method" + # error: [non-subscriptable] "Cannot subscript object of type `NotSubscriptable1` with no `__getitem__` method" union[0] ``` diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/type_guards.md b/crates/ty_python_semantic/resources/mdtest/narrow/type_guards.md index e998f00966..eaefacacdb 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/type_guards.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/type_guards.md @@ -215,7 +215,7 @@ def _(a: tuple[str, int] | tuple[int, str], c: C[Any]): # TODO: Should be `tuple[int, str]` reveal_type(a) # revealed: tuple[str, int] | tuple[int, str] # TODO: Should be `str` - reveal_type(a[1]) # revealed: str | int + reveal_type(a[1]) # revealed: int | str if reveal_type(is_int(a[0])): # revealed: TypeIs[int @ a[0]] # TODO: Should be `tuple[int, str]` diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/class.md b/crates/ty_python_semantic/resources/mdtest/subscript/class.md index 8edb2d94e0..14d82f43f5 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/class.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/class.md @@ -63,12 +63,12 @@ def _(flag: bool): else: class Spam: ... - # error: [possibly-unbound-implicit-call] "Method `__class_getitem__` of type ` | ` is possibly unbound" - # revealed: str + # error: [non-subscriptable] "Cannot subscript object of type `` with no `__class_getitem__` method" + # revealed: str | Unknown reveal_type(Spam[42]) ``` -## TODO: Class getitem non-class union +## Class getitem non-class union ```py def _(flag: bool): @@ -80,8 +80,7 @@ def _(flag: bool): else: Eggs = 1 - a = Eggs[42] # error: "Cannot subscript object of type ` | Literal[1]` with no `__getitem__` method" + a = Eggs[42] # error: "Cannot subscript object of type `Literal[1]` with no `__getitem__` method" - # TODO: should _probably_ emit `str | Unknown` - reveal_type(a) # revealed: Unknown + reveal_type(a) # revealed: str | Unknown ``` diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md b/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md index 49ef6f6f37..9b4edb75c1 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md @@ -188,3 +188,30 @@ class C(Tuple): ... # revealed: tuple[, , , , , , , typing.Protocol, typing.Generic, ] reveal_type(C.__mro__) ``` + +### Union subscript access + +```py +def test(val: tuple[str] | tuple[int]): + reveal_type(val[0]) # revealed: str | int + +def test2(val: tuple[str, None] | list[int | float]): + reveal_type(val[0]) # revealed: str | int | float +``` + +### Union subscript access with non-indexable type + +```py +def test3(val: tuple[str] | tuple[int] | int): + # error: [non-subscriptable] "Cannot subscript object of type `int` with no `__getitem__` method" + reveal_type(val[0]) # revealed: str | int | Unknown +``` + +### Intersection subscript access + +```py +from ty_extensions import Intersection, Not + +def test4(val: Intersection[tuple[str], tuple[int]]): + reveal_type(val[0]) # revealed: str & int +``` diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index e11a31ab42..7ceec5fbbc 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -8125,7 +8125,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } fn infer_subscript_expression_types( - &mut self, + &self, value_node: &ast::Expr, value_ty: Type<'db>, slice_ty: Type<'db>, @@ -8140,7 +8140,22 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { slice_ty, ) } - + // If the value type is a union make sure to union the load types. + // For example: + // val: tuple[int] | tuple[str] + // val[0] can be an int or str type + (Type::Union(union_ty), _, _) => union_ty.map(self.db(), |ty| { + self.infer_subscript_expression_types(value_node, *ty, slice_ty) + }), + (Type::Intersection(intersection_ty), _, _) => intersection_ty + .positive(self.db()) + .iter() + .map(|ty| self.infer_subscript_expression_types(value_node, *ty, slice_ty)) + .fold( + IntersectionBuilder::new(self.db()), + IntersectionBuilder::add_positive, + ) + .build(), // Ex) Given `("a", "b", "c", "d")[1]`, return `"b"` (Type::Tuple(tuple_ty), Type::IntLiteral(int), _) if i32::try_from(int).is_ok() => { let tuple = tuple_ty.tuple(self.db()); @@ -8446,25 +8461,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ); } - match value_ty { - Type::ClassLiteral(_) => { - // TODO: proper support for generic classes - // For now, just infer `Sequence`, if we see something like `Sequence[str]`. This allows us - // to look up attributes on generic base classes, even if we don't understand generics yet. - // Note that this isn't handled by the clause up above for generic classes - // that use legacy type variables and an explicit `Generic` base class. - // Once we handle legacy typevars, this special case will be removed in - // favor of the specialization logic above. - value_ty - } - _ => Type::unknown(), - } + Type::unknown() } } } fn legacy_generic_class_context( - &mut self, + &self, value_node: &ast::Expr, typevars: &[Type<'db>], origin: LegacyGenericBase,