[ty] Narrow on negative subscript indexing (#22682)

## Summary

Negative subscripts are also indicative of a thing being subcriptable:

```python
class Subscriptable:
    def __getitem__(self, key: int) -> int:
        return 42

class NotSubscriptable: ...

def _(x: list[Subscriptable | NotSubscriptable]):
    if not isinstance(x[-1], NotSubscriptable):
        # After narrowing, x[-1] excludes NotSubscriptable, which means subscripting works
        reveal_type(x[-1])  # revealed: Subscriptable & ~NotSubscriptable
        reveal_type(x[-1][0])  # revealed: int
```
This commit is contained in:
Charlie Marsh
2026-01-19 09:04:08 -05:00
committed by GitHub
parent 0793bfdb16
commit 2abaff046e
2 changed files with 118 additions and 0 deletions

View File

@@ -398,3 +398,84 @@ d = D()
if d.c is not None and d.c[0].x[0] is not None:
reveal_type(d.c[0].x[0]) # revealed: int
```
## Narrowing with negative subscripts
Narrowing should work with negative subscripts like `x[-1]`:
```py
def _(x: list[int | None]):
if x[-1] is not None:
reveal_type(x[-1]) # revealed: int
def _(x: list[str | None]):
if x[-1] is None:
reveal_type(x[-1]) # revealed: None
else:
reveal_type(x[-1]) # revealed: str
```
Nested negative subscripts should also work:
```py
def _(x: list[list[int | None]]):
if x[-1][-1] is not None:
reveal_type(x[-1][-1]) # revealed: int
```
Mixed positive and negative subscripts:
```py
def _(x: list[list[int | None]]):
if x[0][-1] is not None:
reveal_type(x[0][-1]) # revealed: int
if x[-1][0] is not None:
reveal_type(x[-1][0]) # revealed: int
```
Attribute access combined with negative subscripts:
```py
class Container:
items: list[int | None]
def _(c: Container):
if c.items[-1] is not None:
reveal_type(c.items[-1]) # revealed: int
```
Multiple conditions in an `and` chain:
```py
def _(x: list[int | None]):
# Narrowing should persist through `and` chains
if x[-1] is not None and x[-1] > 0:
reveal_type(x[-1]) # revealed: int
```
Negative indices with tuples:
```py
def _(t: tuple[int, str, None] | tuple[None, None, int]):
if t[-1] is not None:
reveal_type(t) # revealed: tuple[None, None, int]
else:
reveal_type(t) # revealed: tuple[int, str, None]
if t[-3] is not None:
reveal_type(t) # revealed: tuple[int, str, None]
```
## Narrowing with explicit positive subscripts
Narrowing should work with explicit positive subscripts like `x[+1]`:
```py
def _(x: list[int | None]):
if x[+0] is not None:
reveal_type(x[+0]) # revealed: int
if x[+1] is not None:
reveal_type(x[+1]) # revealed: int
```

View File

@@ -184,6 +184,7 @@ impl MemberExpr {
let start_offset = path.text_len();
match &*subscript.slice {
// Handle integer subscripts, like `x[0]`.
ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
value: ast::Number::Int(index),
..
@@ -192,6 +193,42 @@ impl MemberExpr {
segments
.push(SegmentInfo::new(SegmentKind::IntSubscript, start_offset));
}
// Handle negative integer subscripts, like `x[-1]`.
ast::Expr::UnaryOp(ast::ExprUnaryOp {
op: ast::UnaryOp::USub,
operand,
..
}) => match operand.as_ref() {
ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
value: ast::Number::Int(index),
..
}) => {
let _ = write!(path, "-{index}");
segments.push(SegmentInfo::new(
SegmentKind::IntSubscript,
start_offset,
));
}
_ => return None,
},
// Handle positive integer subscripts with explicit plus, like `x[+1]`.
ast::Expr::UnaryOp(ast::ExprUnaryOp {
op: ast::UnaryOp::UAdd,
operand,
..
}) => match operand.as_ref() {
ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
value: ast::Number::Int(index),
..
}) => {
let _ = write!(path, "{index}");
segments.push(SegmentInfo::new(
SegmentKind::IntSubscript,
start_offset,
));
}
_ => return None,
},
ast::Expr::StringLiteral(string) => {
let _ = write!(path, "{}", string.value);
segments