Skip to content

Commit

Permalink
[flake8-pyi] Add autofix for PYI058 (#9355)
Browse files Browse the repository at this point in the history
## Summary

This PR adds an autofix for the newly added PYI058 rule (added in
#9313). ~~The PR's current implementation is that the fix is only
available if the fully qualified name of `Generator` or `AsyncGenerator`
is being used:~~
- ~~`-> typing.Generator` is converted to `-> typing.Iterator`;~~
- ~~`-> collections.abc.AsyncGenerator[str, Any]` is converted to `->
collections.abc.AsyncIterator[str]`;~~
- ~~but `-> Generator` is _not_ converted to `-> Iterator`. (It would
require more work to figure out if `Iterator` was already imported or
not. And if it wasn't, where should we import it from? `typing`,
`typing_extensions`, or `collections.abc`? It seems much more
complicated.)~~

The fix is marked as always safe for `__iter__` or `__aiter__` methods
in `.pyi` files, but unsafe for all such methods in `.py` files that
have more than one statement in the method body.

This felt slightly fiddly to accomplish, but I couldn't _see_ any
utilities in
https://github.com/astral-sh/ruff/tree/main/crates/ruff_linter/src/fix
that would have made it simpler to implement. Lmk if I'm missing
something, though -- my first time implementing an autofix! :)

## Test Plan

`cargo test` / `cargo insta review`.
  • Loading branch information
AlexWaygood authored Jan 3, 2024
1 parent dc5094d commit 1ffc738
Show file tree
Hide file tree
Showing 5 changed files with 845 additions and 226 deletions.
214 changes: 153 additions & 61 deletions crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI058.py
Original file line number Diff line number Diff line change
@@ -1,82 +1,174 @@
import collections.abc
import typing
from collections.abc import AsyncGenerator, Generator
from typing import Any
def scope():
from collections.abc import Generator

class IteratorReturningSimpleGenerator1:
def __iter__(self) -> Generator: # PYI058 (use `Iterator`)
return (x for x in range(42))
class IteratorReturningSimpleGenerator1:
def __iter__(self) -> Generator:
... # PYI058 (use `Iterator`)

class IteratorReturningSimpleGenerator2:
def __iter__(self, /) -> collections.abc.Generator[str, Any, None]: # PYI058 (use `Iterator`)
"""Fully documented, because I'm a runtime function!"""
yield from "abcdefg"
return None

class IteratorReturningSimpleGenerator3:
def __iter__(self, /) -> collections.abc.Generator[str, None, typing.Any]: # PYI058 (use `Iterator`)
yield "a"
yield "b"
yield "c"
return
def scope():
import typing

class AsyncIteratorReturningSimpleAsyncGenerator1:
def __aiter__(self) -> typing.AsyncGenerator: pass # PYI058 (Use `AsyncIterator`)
class IteratorReturningSimpleGenerator2:
def __iter__(self) -> typing.Generator:
... # PYI058 (use `Iterator`)

class AsyncIteratorReturningSimpleAsyncGenerator2:
def __aiter__(self, /) -> collections.abc.AsyncGenerator[str, Any]: ... # PYI058 (Use `AsyncIterator`)

class AsyncIteratorReturningSimpleAsyncGenerator3:
def __aiter__(self, /) -> collections.abc.AsyncGenerator[str, None]: pass # PYI058 (Use `AsyncIterator`)
def scope():
import collections.abc

class CorrectIterator:
def __iter__(self) -> Iterator[str]: ... # OK
class IteratorReturningSimpleGenerator3:
def __iter__(self) -> collections.abc.Generator:
... # PYI058 (use `Iterator`)

class CorrectAsyncIterator:
def __aiter__(self) -> collections.abc.AsyncIterator[int]: ... # OK

class Fine:
def __iter__(self) -> typing.Self: ... # OK
def scope():
import collections.abc
from typing import Any

class StrangeButWeWontComplainHere:
def __aiter__(self) -> list[bytes]: ... # OK
class IteratorReturningSimpleGenerator4:
def __iter__(self, /) -> collections.abc.Generator[str, Any, None]:
... # PYI058 (use `Iterator`)

def __iter__(self) -> Generator: ... # OK (not in class scope)
def __aiter__(self) -> AsyncGenerator: ... # OK (not in class scope)

class IteratorReturningComplexGenerator:
def __iter__(self) -> Generator[str, int, bytes]: ... # OK
def scope():
import collections.abc
import typing

class AsyncIteratorReturningComplexAsyncGenerator:
def __aiter__(self) -> AsyncGenerator[str, int]: ... # OK
class IteratorReturningSimpleGenerator5:
def __iter__(self, /) -> collections.abc.Generator[str, None, typing.Any]:
... # PYI058 (use `Iterator`)

class ClassWithInvalidAsyncAiterMethod:
async def __aiter__(self) -> AsyncGenerator: ... # OK

class IteratorWithUnusualParameters1:
def __iter__(self, foo) -> Generator: ... # OK
def scope():
from collections.abc import Generator

class IteratorWithUnusualParameters2:
def __iter__(self, *, bar) -> Generator: ... # OK
class IteratorReturningSimpleGenerator6:
def __iter__(self, /) -> Generator[str, None, None]:
... # PYI058 (use `Iterator`)

class IteratorWithUnusualParameters3:
def __iter__(self, *args) -> Generator: ... # OK

class IteratorWithUnusualParameters4:
def __iter__(self, **kwargs) -> Generator: ... # OK
def scope():
import typing_extensions

class IteratorWithIterMethodThatReturnsThings:
def __iter__(self) -> Generator: # OK
yield
return 42
class AsyncIteratorReturningSimpleAsyncGenerator1:
def __aiter__(
self,
) -> typing_extensions.AsyncGenerator:
... # PYI058 (Use `AsyncIterator`)

class IteratorWithIterMethodThatReceivesThingsFromSend:
def __iter__(self) -> Generator: # OK
x = yield 42

class IteratorWithNonTrivialIterBody:
def __iter__(self) -> Generator: # OK
foo, bar, baz = (1, 2, 3)
yield foo
yield bar
yield baz
def scope():
import collections.abc

class AsyncIteratorReturningSimpleAsyncGenerator2:
def __aiter__(self, /) -> collections.abc.AsyncGenerator[str, Any]:
... # PYI058 (Use `AsyncIterator`)


def scope():
import collections.abc

class AsyncIteratorReturningSimpleAsyncGenerator3:
def __aiter__(self, /) -> collections.abc.AsyncGenerator[str, None]:
... # PYI058 (Use `AsyncIterator`)


def scope():
from typing import Iterator

class CorrectIterator:
def __iter__(self) -> Iterator[str]:
... # OK


def scope():
import collections.abc

class CorrectAsyncIterator:
def __aiter__(self) -> collections.abc.AsyncIterator[int]:
... # OK


def scope():
import typing

class Fine:
def __iter__(self) -> typing.Self:
... # OK


def scope():
class StrangeButWeWontComplainHere:
def __aiter__(self) -> list[bytes]:
... # OK


def scope():
from collections.abc import Generator

def __iter__(self) -> Generator:
... # OK (not in class scope)


def scope():
from collections.abc import AsyncGenerator

def __aiter__(self) -> AsyncGenerator:
... # OK (not in class scope)


def scope():
from collections.abc import Generator

class IteratorReturningComplexGenerator:
def __iter__(self) -> Generator[str, int, bytes]:
... # OK


def scope():
from collections.abc import AsyncGenerator

class AsyncIteratorReturningComplexAsyncGenerator:
def __aiter__(self) -> AsyncGenerator[str, int]:
... # OK


def scope():
from collections.abc import AsyncGenerator

class ClassWithInvalidAsyncAiterMethod:
async def __aiter__(self) -> AsyncGenerator:
... # OK


def scope():
from collections.abc import Generator

class IteratorWithUnusualParameters1:
def __iter__(self, foo) -> Generator:
... # OK


def scope():
from collections.abc import Generator

class IteratorWithUnusualParameters2:
def __iter__(self, *, bar) -> Generator:
... # OK


def scope():
from collections.abc import Generator

class IteratorWithUnusualParameters3:
def __iter__(self, *args) -> Generator:
... # OK


def scope():
from collections.abc import Generator

class IteratorWithUnusualParameters4:
def __iter__(self, **kwargs) -> Generator:
... # OK
Loading

0 comments on commit 1ffc738

Please sign in to comment.