From 5a18e93d658718efdaa2ae849befd92a729be03e Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Thu, 15 Jan 2026 00:18:54 +0000 Subject: [PATCH] [ty] Make special cases for subscript inference exhaustive (#22035) ## Summary Fixes https://github.com/astral-sh/ty/issues/2015. We weren't recursing into the value of a type alias when we should have been. There are situations where we should also be recursing into the bounds/constraints of a typevar. I initially tried to do that as well in this PR, but that seems... trickier. For now I'm cutting scope; this PR does, however, add several failing tests for those cases. ## Test Plan added mdtests --- .../resources/mdtest/subscript/alias.md | 43 +++++++++ .../resources/mdtest/subscript/class.md | 2 +- .../resources/mdtest/subscript/string.md | 11 +++ .../resources/mdtest/subscript/tuple.md | 2 +- .../resources/mdtest/subscript/typevar.md | 89 +++++++++++++++++++ .../src/types/infer/builder.rs | 68 +++++++++++++- 6 files changed, 210 insertions(+), 5 deletions(-) create mode 100644 crates/ty_python_semantic/resources/mdtest/subscript/alias.md create mode 100644 crates/ty_python_semantic/resources/mdtest/subscript/typevar.md diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/alias.md b/crates/ty_python_semantic/resources/mdtest/subscript/alias.md new file mode 100644 index 0000000000..cd0485f9b9 --- /dev/null +++ b/crates/ty_python_semantic/resources/mdtest/subscript/alias.md @@ -0,0 +1,43 @@ +# Subscripts involving type aliases + +Aliases are expanded during analysis of subscripts. + +```toml +[environment] +python-version = "3.12" +``` + +```py +from typing_extensions import TypeAlias, Literal + +ImplicitTuple = tuple[str, int, int] +PEP613Tuple: TypeAlias = tuple[str, int, int] +type PEP695Tuple = tuple[str, int, int] + +ImplicitZero = Literal[0] +PEP613Zero: TypeAlias = Literal[0] +type PEP695Zero = Literal[0] + +def f( + implicit_tuple: ImplicitTuple, + pep_613_tuple: PEP613Tuple, + pep_695_tuple: PEP695Tuple, + implicit_zero: ImplicitZero, + pep_613_zero: PEP613Zero, + pep_695_zero: PEP695Zero, +): + reveal_type(implicit_tuple[:2]) # revealed: tuple[str, int] + reveal_type(implicit_tuple[implicit_zero]) # revealed: str + reveal_type(implicit_tuple[pep_613_zero]) # revealed: str + reveal_type(implicit_tuple[pep_695_zero]) # revealed: str + + reveal_type(pep_613_tuple[:2]) # revealed: tuple[str, int] + reveal_type(pep_613_tuple[implicit_zero]) # revealed: str + reveal_type(pep_613_tuple[pep_613_zero]) # revealed: str + reveal_type(pep_613_tuple[pep_695_zero]) # revealed: str + + reveal_type(pep_695_tuple[:2]) # revealed: tuple[str, int] + reveal_type(pep_695_tuple[implicit_zero]) # revealed: str + reveal_type(pep_695_tuple[pep_613_zero]) # revealed: str + reveal_type(pep_695_tuple[pep_695_zero]) # revealed: str +``` diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/class.md b/crates/ty_python_semantic/resources/mdtest/subscript/class.md index 1c14dd23a5..0b680e5310 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/class.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/class.md @@ -106,5 +106,5 @@ class Bar: def f(x: Foo): if isinstance(x, Bar): # TODO: should be `int` - reveal_type(x["whatever"]) # revealed: @Todo(Subscript expressions on intersections) + reveal_type(x["whatever"]) # revealed: @Todo(Subscript expressions with intersections) ``` diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/string.md b/crates/ty_python_semantic/resources/mdtest/subscript/string.md index 469300c0b4..7e8e59f504 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/string.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/string.md @@ -80,6 +80,17 @@ def _(m: int, n: int, s2: str): reveal_type(substring2) # revealed: str ``` +## LiteralString + +```py +from typing_extensions import LiteralString + +def f(x: LiteralString): + reveal_type(x[0]) # revealed: LiteralString + reveal_type(x[True]) # revealed: LiteralString + reveal_type(x[1:42]) # revealed: LiteralString +``` + ## Unsupported slice types ```py diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md b/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md index 1ec85b3203..8f686249fb 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/tuple.md @@ -430,5 +430,5 @@ class Bar: ... def test4(val: Intersection[tuple[Foo], tuple[Bar]]): # TODO: should be `Foo & Bar` - reveal_type(val[0]) # revealed: @Todo(Subscript expressions on intersections) + reveal_type(val[0]) # revealed: @Todo(Subscript expressions with intersections) ``` diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/typevar.md b/crates/ty_python_semantic/resources/mdtest/subscript/typevar.md new file mode 100644 index 0000000000..e2917c141c --- /dev/null +++ b/crates/ty_python_semantic/resources/mdtest/subscript/typevar.md @@ -0,0 +1,89 @@ +# Subscripts involving type variables + +## TypeVar bound/constrained to a tuple/int-literal/bool-literal + +The upper bounds of type variables are considered when analysing subscripts. + +```toml +[environment] +python-version = "3.12" +``` + +```py +from typing_extensions import TypeAlias, Literal + +ImplicitTuple = tuple[str, int, int] +PEP613Tuple: TypeAlias = tuple[str, int, int] +type PEP695Tuple = tuple[str, int, int] + +ImplicitZero = Literal[0] +PEP613Zero: TypeAlias = Literal[0] +type PEP695Zero = Literal[0] + +# fmt: off + +def f[ + BoundedTupleT: tuple[str, int, bytes], + ConstrainedTupleT: (tuple[str, int, bytes], tuple[int, bytes, str]), + BoundedZeroT: Literal[0], + ConstrainedIntLiteralT: (Literal[0], Literal[1]) +]( + tuple_1: BoundedTupleT, + tuple_2: ConstrainedTupleT, + zero: BoundedZeroT, + some_integer: ConstrainedIntLiteralT, +): + # TODO: would ideally be `tuple[str, int]` + reveal_type(tuple_1[:2]) # revealed: tuple[str | int | bytes, ...] + reveal_type(tuple_1[zero]) # revealed: str + + # TODO: ideally this might be `str | int`, + # but it's hard to do that without introducing false positives elsewhere + reveal_type(tuple_1[some_integer]) # revealed: str | int | bytes + + # TODO: would ideally be `tuple[str, int] | tuple[int, bytes]` + reveal_type(tuple_2[:2]) # revealed: tuple[str | int | bytes, ...] + reveal_type(tuple_2[zero]) # revealed: str | int + reveal_type(tuple_2[some_integer]) # revealed: str | int | bytes + +# fmt: on +``` + +## TypeVars + +```toml +[environment] +python-version = "3.12" +``` + +```py +from typing import Protocol + +class SupportsLessThan(Protocol): + def __lt__(self, other, /) -> bool: ... + +def f[K: SupportsLessThan](dictionary: dict[K, int], key: K): + reveal_type(dictionary[key]) # revealed: int +``` + +## ParamSpecs + +```toml +[environment] +python-version = "3.12" +``` + +```py +from typing import Callable + +def decorator[**P, T](func: Callable[P, T]) -> Callable[P, T]: + def inner(*args: P.args, **kwargs: P.kwargs) -> T: + if len(args) > 0: + # error: [invalid-assignment] + args = args[1:] + + # `func` requires the full `ParamSpec` passed into `decorator`, + # but here the first argument is skipped, so we should possibly emit an error here: + return func(*args, **kwargs) + return inner +``` diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 351c8cc3a1..97b85d94b8 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -14091,16 +14091,36 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let value_node = subscript.value.as_ref(); let inferred = match (value_ty, slice_ty) { + (Type::Dynamic(_) | Type::Never, _) => Some(value_ty), + + (Type::TypeAlias(alias), _) => Some(self.infer_subscript_expression_types( + subscript, + alias.value_type(self.db()), + slice_ty, + expr_context, + )), + + (_, Type::TypeAlias(alias)) => Some(self.infer_subscript_expression_types( + subscript, + value_ty, + alias.value_type(self.db()), + expr_context, + )), + (Type::Union(union), _) => Some(union.map(db, |element| { self.infer_subscript_expression_types(subscript, *element, slice_ty, expr_context) })), + (_, Type::Union(union)) => Some(union.map(db, |element| { + self.infer_subscript_expression_types(subscript, value_ty, *element, expr_context) + })), + // TODO: we can map over the intersection and fold the results back into an intersection, // but we need to make sure we avoid emitting a diagnostic if one positive element has a `__getitem__` // method but another does not. This means `infer_subscript_expression_types` // needs to return a `Result` rather than eagerly emitting diagnostics. - (Type::Intersection(_), _) => { - Some(todo_type!("Subscript expressions on intersections")) + (Type::Intersection(_), _) | (_, Type::Intersection(_)) => { + Some(todo_type!("Subscript expressions with intersections")) } // Ex) Given `("a", "b", "c", "d")[1]`, return `"b"` @@ -14180,6 +14200,16 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } }), + (Type::LiteralString, Type::IntLiteral(_) | Type::BooleanLiteral(_)) => { + Some(Type::LiteralString) + } + + (Type::LiteralString, Type::NominalInstance(nominal)) + if nominal.slice_literal(db).is_some() => + { + Some(Type::LiteralString) + } + // Ex) Given `b"value"[1]`, return `97` (i.e., `ord(b"a")`) (Type::BytesLiteral(literal_ty), Type::IntLiteral(i64_int)) => { i32::try_from(i64_int).ok().map(|i32_int| { @@ -14311,7 +14341,39 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { Some(todo_type!("Inference of subscript on special form")) } - _ => None, + ( + Type::FunctionLiteral(_) + | Type::WrapperDescriptor(_) + | Type::BoundMethod(_) + | Type::DataclassDecorator(_) + | Type::DataclassTransformer(_) + | Type::Callable(_) + | Type::ModuleLiteral(_) + | Type::ClassLiteral(_) + | Type::GenericAlias(_) + | Type::SubclassOf(_) + | Type::AlwaysFalsy + | Type::AlwaysTruthy + | Type::IntLiteral(_) + | Type::BooleanLiteral(_) + | Type::ProtocolInstance(_) + | Type::PropertyInstance(_) + | Type::EnumLiteral(_) + | Type::BoundSuper(_) + | Type::TypeIs(_) + | Type::TypeGuard(_) + | Type::TypedDict(_) + | Type::NewTypeInstance(_) + | Type::NominalInstance(_) + | Type::SpecialForm(_) + | Type::KnownInstance(_) + | Type::StringLiteral(_) + | Type::BytesLiteral(_) + | Type::LiteralString + | Type::TypeVar(_) // TODO: more complex logic required here! + | Type::KnownBoundMethod(_), + _, + ) => None, }; if let Some(inferred) = inferred {