diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts index 68b3343e6fae..bf70a60d041c 100644 --- a/packages/pyright-internal/src/analyzer/typeEvaluator.ts +++ b/packages/pyright-internal/src/analyzer/typeEvaluator.ts @@ -3402,9 +3402,22 @@ export function createTypeEvaluator( const iterReturnType = getTypeOfMagicMethodCall(subtype, iterMethodName, [], errorNode)?.type; if (!iterReturnType) { - // There was no __iter__. See if we can fall back to - // the __getitem__ method instead. + // There was no callable __iter__. The legacy sequence protocol + // falls back to __getitem__, but CPython does that only when + // __iter__ is missing. An explicit `__iter__ = None` (or any + // non-callable __iter__) makes the object not iterable. + let hasExplicitIterMember = false; if (!isAsync && isClassInstance(subtype)) { + hasExplicitIterMember = !!lookUpObjectMember( + subtype, + iterMethodName, + MemberAccessFlags.SkipInstanceMembers | + MemberAccessFlags.SkipAttributeAccessOverride | + MemberAccessFlags.SkipObjectBaseClass + ); + } + + if (!isAsync && isClassInstance(subtype) && !hasExplicitIterMember) { const getItemReturnType = getTypeOfMagicMethodCall( subtype, '__getitem__', diff --git a/packages/pyright-internal/src/tests/samples/forLoop3.py b/packages/pyright-internal/src/tests/samples/forLoop3.py new file mode 100644 index 000000000000..c625c68613db --- /dev/null +++ b/packages/pyright-internal/src/tests/samples/forLoop3.py @@ -0,0 +1,44 @@ +# This sample tests that `__iter__ = None` disables iteration even +# when `__getitem__` is defined. CPython raises TypeError in this case. + + +class SequenceProtocol: + def __getitem__(self, item: int) -> int: + if item >= 3: + raise IndexError + return item + + +# The legacy sequence protocol should still be iterable. +for _ in SequenceProtocol(): + pass + + +class IterNone: + def __getitem__(self, item: int) -> int: + return item + + __iter__ = None + + +# This should generate an error because __iter__ is None. +for _ in IterNone(): + pass + + +class IterNoneSubclass(IterNone): + pass + + +# This should generate an error because __iter__ is None. +for _ in IterNoneSubclass(): + pass + + +class IterRestored(IterNone): + def __iter__(self): + yield 1 + + +for _ in IterRestored(): + pass diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts index 3a5a53469bf7..92b2a5212ad1 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts +++ b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts @@ -518,6 +518,12 @@ test('ForLoop2', () => { TestUtils.validateResults(analysisResults, 7); }); +test('ForLoop3', () => { + const analysisResults = TestUtils.typeAnalyzeSampleFiles(['forLoop3.py']); + + TestUtils.validateResults(analysisResults, 2); +}); + test('Comprehension1', () => { const analysisResults = TestUtils.typeAnalyzeSampleFiles(['comprehension1.py']);