[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
This commit is contained in:
Alex Waygood
2026-01-15 00:18:54 +00:00
committed by GitHub
parent 5d99ef9d95
commit 5a18e93d65
6 changed files with 210 additions and 5 deletions

View File

@@ -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
```

View File

@@ -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)
```

View File

@@ -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

View File

@@ -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)
```

View File

@@ -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
```

View File

@@ -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 {