diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md index b1a324e557..f97bc4b0f5 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md @@ -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 +``` diff --git a/crates/ty_python_semantic/src/semantic_index/member.rs b/crates/ty_python_semantic/src/semantic_index/member.rs index 8511c7b295..6979ac5925 100644 --- a/crates/ty_python_semantic/src/semantic_index/member.rs +++ b/crates/ty_python_semantic/src/semantic_index/member.rs @@ -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