Skip to content

[red-knot] Add type inference for basic for loops#13195

Merged
AlexWaygood merged 5 commits intomainfrom
alex/for-loop-types
Sep 4, 2024
Merged

[red-knot] Add type inference for basic for loops#13195
AlexWaygood merged 5 commits intomainfrom
alex/for-loop-types

Conversation

@AlexWaygood
Copy link
Copy Markdown
Member

@AlexWaygood AlexWaygood commented Sep 1, 2024

Summary

This PR adds type inference for basic for loop variables to red-knot.

Inferring the type of a loop var involves either the new-style iteration protocol (which invokes __iter__ to get an iterator, and then __next__ on that iterator), or the old-style iteration protocol (which passes incrementally larger ints to a __getitem__ method). If __iter__ is defined on the class, __iter__ always takes precedence, even if __iter__ is e.g. set to None: this is a technique that can be used to make classes that define __getitem__ not-iterable.

Whichever protocol is used to make a class's instances iterable, the dunder methods are always looked up on the class of an instance rather than the instance itself. I.e., the interpreter won't consider foo in this snippet to be iterable, because __iter__ only exists on the instance, not the class:

>>> class Foo: ...
... 
>>> foo = Foo()
>>> foo.__iter__ = lambda: iter(range(42))
>>> list(foo.__iter__())
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41]
>>> for x in foo:
...     pass
... 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'Foo' object is not iterable

To emulate this runtime behaviour, I added a Type::to_type_of_class method that allows us to go from a type representing <instance of int> to a type representing <the int class itself>.

Test Plan

cargo test

@AlexWaygood AlexWaygood added the ty Multi-file analysis & type inference label Sep 1, 2024
@github-actions
Copy link
Copy Markdown
Contributor

github-actions bot commented Sep 1, 2024

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

@AlexWaygood AlexWaygood force-pushed the alex/for-loop-types branch 2 times, most recently from 9fdeb3e to 52941b9 Compare September 2, 2024 13:21
@MichaReiser
Copy link
Copy Markdown
Member

All changes look good to me but I prefer for @carljm to take a look

Copy link
Copy Markdown
Contributor

@carljm carljm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is great!! So neat to see things coming together such that we can now do real inference based on dunder protocols :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ty Multi-file analysis & type inference

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants