[flake8-pyi] Add autofix for PYI058 #9355
Conversation
|
crates/ruff_linter/src/rules/flake8_pyi/rules/bad_generator_return_type.rs
Show resolved
Hide resolved
crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI058.py
Outdated
Show resolved
Hide resolved
| Typing, | ||
| TypingExtensions, | ||
| CollectionsAbc, | ||
| } |
There was a problem hiding this comment.
@AlexWaygood - I did some light refactoring here to use enums for all the different values rather than strings. You might be interested in taking a look. I think it helps enforce some of the invariants that are expected throughout the code and make it clear which values are valid in various places.
There was a problem hiding this comment.
Yeah, this is nice, thanks! Fewer magic strings everywhere
| from collections.abc import AsyncGenerator, Generator | ||
| from typing import Any | ||
| def scope(): | ||
| from collections.abc import Generator |
There was a problem hiding this comment.
Unfortunately, I had to scope each test here into its own function so that they can reason about the imports independently.
There was a problem hiding this comment.
Ahh, that makes sense, was wondering if that would be an issue
| checker.generator().expr(yield_type_info.expr), | ||
| yield_type_info.range(), | ||
| ) | ||
| }); |
There was a problem hiding this comment.
@AlexWaygood - I tried to generalize this so that it doesn't need different logic for attribute vs. name.
There was a problem hiding this comment.
Nice, I didn't realise how powerful get_or_import_symbol was
| 29 30 | | ||
| 30 31 | class IteratorReturningSimpleGenerator5: | ||
| 31 |- def __iter__(self, /) -> collections.abc.Generator[str, None, typing.Any]: ... # PYI058 (use `Iterator`) | ||
| 32 |+ def __iter__(self, /) -> Iterator[str]: ... # PYI058 (use `Iterator`) |
There was a problem hiding this comment.
This is a bug in get_or_import_symbol (it doesn't know that it can reuse collections.abc because it's a two-part module). Need to fix.
There was a problem hiding this comment.
I'll fix this separately, since the current fix isn't "wrong", it's just suboptimal.
|
|
||
| #[derive(Debug)] | ||
| struct YieldTypeInfo<'a> { | ||
| expr: &'a ast::Expr, |
There was a problem hiding this comment.
@AlexWaygood - A bit advanced but I changed this to use a lifetime so that it doesn't need to clone the expression. We know the expression will "live" long enough, so it's okay for us to use a reference.
There was a problem hiding this comment.
Nice, thanks! I've read about things being generic over lifetimes but have yet to successfully use the feature :)
| "__iter__" => "Iterator", | ||
| "__aiter__" => "AsyncIterator", | ||
| _ => return, | ||
| }; |
There was a problem hiding this comment.
@AlexWaygood - I removed this since it's enforced in let (method, module, member) = { below.
| better_return_type: String, | ||
| method_name: String, | ||
| return_type: Iterator, | ||
| method: Method, |
There was a problem hiding this comment.
@AlexWaygood - I changed this to use the enums, which implement std::fmt::Display, so we can still inline them in the format! call and automatically get the expected formatting.
charliermarsh
added
fixes
| impl Ranged for YieldTypeInfo<'_> { | ||
| fn range(&self) -> TextRange { | ||
| self.range | ||
| } | ||
| } |
There was a problem hiding this comment.
For my education: what's the advantage of implementing the trait here rather than just reading the attribute?
There was a problem hiding this comment.
It's really minor, more of a consistency thing... If you implement Ranged, then you get methods like .start() and .end() instead of having to do .range.start().
|
Thanks! |
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 ofGeneratororAsyncGeneratoris being used:-> typing.Generatoris converted to-> typing.Iterator;-> collections.abc.AsyncGenerator[str, Any]is converted to-> collections.abc.AsyncIterator[str];but-> Generatoris not converted to-> Iterator. (It would require more work to figure out ifIteratorwas already imported or not. And if it wasn't, where should we import it from?typing,typing_extensions, orcollections.abc? It seems much more complicated.)The fix is marked as always safe for
__iter__or__aiter__methods in.pyifiles, but unsafe for all such methods in.pyfiles 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.