From 1de9af0cf48992c5d936229e3110bed7b8027d49 Mon Sep 17 00:00:00 2001 From: Kyle Into Date: Sat, 17 Oct 2026 00:28:16 -0700 Subject: [PATCH 1/2] Add test for subscripting a TypeVar bounded by a class whose `__getitem__` returns `Self` Summary: Subscripting a value whose type is a TypeVar bounded by a class (`seq[0:1]` with `S: SupportsSlice`) infers the bound instead of `S` when `__getitem__` returns `Self`. Calling `seq.__getitem__(slice(0, 1))` directly infers `S` as expected. This adds a test, marked `bug`, that records the current behavior. The next diff fixes it. reproduces https://github.com/facebook/pyrefly/issues/5257 Differential Revision: D125829599 --- pyrefly/lib/test/generic_restrictions.rs | 59 ++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/pyrefly/lib/test/generic_restrictions.rs b/pyrefly/lib/test/generic_restrictions.rs index d376c7d37c..88ef58c279 100644 --- a/pyrefly/lib/test/generic_restrictions.rs +++ b/pyrefly/lib/test/generic_restrictions.rs @@ -755,6 +755,65 @@ def foo[T: (A, B)](x: T) -> T: "#, ); +testcase!( + bug = "subscripting a bounded TypeVar resolves `Self` to the bound instead of the TypeVar", + test_bounded_typevar_subscript_return_self, + r#" +from typing import Protocol, Self, assert_type + +class SupportsSlice(Protocol): + def __getitem__(self, index: slice) -> Self: ... + +class Concrete: + def __getitem__(self, index: slice) -> Self: + return self + +def protocol_bound[S: SupportsSlice](seq: S) -> None: + assert_type(seq[0:1], S) # E: assert_type(SupportsSlice, S) failed + +def class_bound[S: Concrete](seq: S) -> None: + assert_type(seq[0:1], S) # E: assert_type(Concrete, S) failed + +def tuple_bound[S: tuple[int, str]](seq: S) -> None: + assert_type(seq[0], int) + "#, +); + +testcase!( + bug = "subscripting a restricted TypeVar resolves `Self` to the restriction", + test_restricted_typevar_subscript_return_self, + r#" +from typing import NamedTuple, Self, assert_type + +class A: + def __getitem__(self, index: slice) -> Self: + return self + +class B: + def __getitem__(self, index: slice) -> Self: + return self + +class NT(NamedTuple): + x: int + y: str + +def union_bound[S: A | B](seq: S) -> None: + assert_type(seq[0:1], S) # E: assert_type(A | B, S) failed + +def constrained[S: (A, B)](seq: S) -> None: + assert_type(seq[0:1], S) # E: assert_type(A | B, S) failed + +def named_tuple_bound[S: NT](seq: S) -> None: + assert_type(seq[0], int) + +def list_bound[S: list[int]](seq: S) -> None: + assert_type(seq[0:1], list[int]) + +def not_subscriptable[S: int](x: S) -> None: + x[0] # E: Cannot index into `int` + "#, +); + testcase!( test_typevar_single_constraint_is_error, r#" From 4e58de28f06523e9226ce46274fe530970c1e4b8 Mon Sep 17 00:00:00 2001 From: Kyle Into Date: Sat, 17 Oct 2026 00:28:16 -0700 Subject: [PATCH 2/2] Resolve `Self` to the TypeVar when subscripting a class-bounded TypeVar Summary: Subscripting a value whose type is a TypeVar bounded by a class (`seq[0:1]` with `S: SupportsSlice`) inferred the bound instead of `S` when `__getitem__` returns `Self`. Calling the method explicitly (`seq.__getitem__(slice(0, 1))`) already inferred `S`. The subscript path handled a bounded TypeVar by subscripting its bound, so `__getitem__` was looked up and called with the bound as the receiver, and `Self` resolved to the bound. Binary operators don't have this problem because they call the dunder on the TypeVar itself. For a TypeVar bounded by a class whose subscripts use the plain `__getitem__` call, we now call `__getitem__` on the TypeVar, the same way an explicit method call does. Attribute lookup on a bounded TypeVar binds `Self` to the TypeVar. Bounds with builtin subscript handling (lists, tuples, shaped arrays, `nn.ModuleDict`) still subscript the bound, so their precise indexing is kept. fixes https://github.com/facebook/pyrefly/issues/5257 Differential Revision: D125829160 --- pyrefly/lib/alt/expr.rs | 67 ++++++++++++++++-------- pyrefly/lib/test/generic_restrictions.rs | 12 ++--- 2 files changed, 51 insertions(+), 28 deletions(-) diff --git a/pyrefly/lib/alt/expr.rs b/pyrefly/lib/alt/expr.rs index 7b29c2cdab..89fa182b33 100644 --- a/pyrefly/lib/alt/expr.rs +++ b/pyrefly/lib/alt/expr.rs @@ -3420,6 +3420,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let mut aliases = SmallSet::new(); self.subscript_infer_for_type_with_key_present_inner( base, + None, slice, range, errors, @@ -3429,9 +3430,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ) } + /// `receiver`, when set, is the restricted TypeVar whose bound or constraint `base` is. + /// `__getitem__` is then called on the TypeVar, so that `Self` in its signature binds to + /// the TypeVar, while `base` still selects any builtin subscript handling. fn subscript_infer_for_type_with_key_present_inner( &self, base: &Type, + receiver: Option<&Type>, slice: &Expr, range: TextRange, errors: &ErrorCollector, @@ -3467,6 +3472,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let mut aliases = aliases.clone(); self.subscript_infer_for_type_with_key_present_inner( branch, + receiver, slice, range, errors, @@ -3760,15 +3766,18 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.infer_shaped_array_index(shaped_array_type, slice, range, errors) } // Shaped arrays that opt out of built-in indexing use their declared method. - Type::ShapedArray(_) => self.call_method_or_error( - &base, - &dunder::GETITEM, - range, - &[CallArg::expr(slice)], - &[], - errors, - Some(&|| ErrorContext::Index(self.for_display(base.clone()))), - ), + Type::ShapedArray(_) => { + let receiver = receiver.unwrap_or(&base); + self.call_method_or_error( + receiver, + &dunder::GETITEM, + range, + &[CallArg::expr(slice)], + &[], + errors, + Some(&|| ErrorContext::Index(self.for_display(receiver.clone()))), + ) + } // Shaped arrays that have not gone through annotation // canonicalization still use tensor indexing logic. Type::ClassType(ref cls) @@ -3792,32 +3801,42 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } // Special handling for nn.ModuleDict with TypedDict type argument Type::ClassType(ref cls) if is_nn_module_dict(cls) => { - self.try_nn_module_dict_index(cls, &base, slice, range, errors) + self.try_nn_module_dict_index( + cls, + receiver.unwrap_or(&base), + slice, + range, + errors, + ) } Type::ClassType(ref cls) | Type::SelfType(ref cls) if cls.is_builtin("list") => { let index_arg = match builtin_sequence_slice_ty { Some(ty) => CallArg::ty(ty, slice.range()), None => CallArg::expr(slice), }; + let receiver = receiver.unwrap_or(&base); self.call_method_or_error( - &base, + receiver, &dunder::GETITEM, range, &[index_arg], &[], errors, - Some(&|| ErrorContext::Index(self.for_display(base.clone()))), + Some(&|| ErrorContext::Index(self.for_display(receiver.clone()))), + ) + } + Type::ClassType(_) | Type::SelfType(_) => { + let receiver = receiver.unwrap_or(&base); + self.call_method_or_error( + receiver, + &dunder::GETITEM, + range, + &[CallArg::expr(slice)], + &[], + errors, + Some(&|| ErrorContext::Index(self.for_display(receiver.clone()))), ) } - Type::ClassType(_) | Type::SelfType(_) => self.call_method_or_error( - &base, - &dunder::GETITEM, - range, - &[CallArg::expr(slice)], - &[], - errors, - Some(&|| ErrorContext::Index(self.for_display(base.clone()))), - ), Type::DataFrame(schema) => { if let Expr::List(ExprList { elts, .. }) = slice && schema.kind == DataFrameKind::Polars @@ -3853,6 +3872,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } let result = self.subscript_infer_for_type_with_key_present_inner( &schema.underlying_type(), + receiver, slice, range, errors, @@ -3876,6 +3896,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } Type::Series(schema) => self.subscript_infer_for_type_with_key_present_inner( &schema.underlying_type(), + receiver, slice, range, errors, @@ -3888,6 +3909,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { Restriction::Bound(bound) => self .subscript_infer_for_type_with_key_present_inner( bound, + Some(&base), slice, range, errors, @@ -3899,6 +3921,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.unions(constraints.map(|constraint| { self.subscript_infer_for_type_with_key_present_inner( constraint, + Some(&base), slice, range, errors, @@ -3911,6 +3934,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { Restriction::ShapeExtension(extension) => self .subscript_infer_for_type_with_key_present_inner( &extension.upper_bound(self.stdlib, self.heap), + Some(&base), slice, range, errors, @@ -4014,6 +4038,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } else { let result = self.subscript_infer_for_type_with_key_present_inner( &self.untype_alias(&ta), + receiver, slice, range, errors, diff --git a/pyrefly/lib/test/generic_restrictions.rs b/pyrefly/lib/test/generic_restrictions.rs index 88ef58c279..a8ce1e6428 100644 --- a/pyrefly/lib/test/generic_restrictions.rs +++ b/pyrefly/lib/test/generic_restrictions.rs @@ -756,7 +756,6 @@ def foo[T: (A, B)](x: T) -> T: ); testcase!( - bug = "subscripting a bounded TypeVar resolves `Self` to the bound instead of the TypeVar", test_bounded_typevar_subscript_return_self, r#" from typing import Protocol, Self, assert_type @@ -769,10 +768,10 @@ class Concrete: return self def protocol_bound[S: SupportsSlice](seq: S) -> None: - assert_type(seq[0:1], S) # E: assert_type(SupportsSlice, S) failed + assert_type(seq[0:1], S) def class_bound[S: Concrete](seq: S) -> None: - assert_type(seq[0:1], S) # E: assert_type(Concrete, S) failed + assert_type(seq[0:1], S) def tuple_bound[S: tuple[int, str]](seq: S) -> None: assert_type(seq[0], int) @@ -780,7 +779,6 @@ def tuple_bound[S: tuple[int, str]](seq: S) -> None: ); testcase!( - bug = "subscripting a restricted TypeVar resolves `Self` to the restriction", test_restricted_typevar_subscript_return_self, r#" from typing import NamedTuple, Self, assert_type @@ -798,10 +796,10 @@ class NT(NamedTuple): y: str def union_bound[S: A | B](seq: S) -> None: - assert_type(seq[0:1], S) # E: assert_type(A | B, S) failed + assert_type(seq[0:1], S) def constrained[S: (A, B)](seq: S) -> None: - assert_type(seq[0:1], S) # E: assert_type(A | B, S) failed + assert_type(seq[0:1], S) def named_tuple_bound[S: NT](seq: S) -> None: assert_type(seq[0], int) @@ -810,7 +808,7 @@ def list_bound[S: list[int]](seq: S) -> None: assert_type(seq[0:1], list[int]) def not_subscriptable[S: int](x: S) -> None: - x[0] # E: Cannot index into `int` + x[0] # E: Cannot index into `S` "#, );