[flake8-pyi] Implement PYI045 (#4700)

crates/ruff/resources/test/fixtures/flake8_pyi/PYI045.py

@@ -0,0 +1,85 @@
+import collections.abc
+import typing
+from collections.abc import Iterator, Iterable
+
+
+class NoReturn:
+    def __iter__(self):
+        ...
+
+
+class TypingIterableTReturn:
+    def __iter__(self) -> typing.Iterable[int]:
+        ...
+
+    def not_iter(self) -> typing.Iterable[int]:
+        ...
+
+
+class TypingIterableReturn:
+    def __iter__(self) -> typing.Iterable:
+        ...
+
+    def not_iter(self) -> typing.Iterable:
+        ...
+
+
+class CollectionsIterableTReturn:
+    def __iter__(self) -> collections.abc.Iterable[int]:
+        ...
+
+    def not_iter(self) -> collections.abc.Iterable[int]:
+        ...
+
+
+class CollectionsIterableReturn:
+    def __iter__(self) -> collections.abc.Iterable:
+        ...
+
+    def not_iter(self) -> collections.abc.Iterable:
+        ...
+
+
+class IterableReturn:
+    def __iter__(self) -> Iterable:
+        ...
+
+
+class IteratorReturn:
+    def __iter__(self) -> Iterator:
+        ...
+
+
+class IteratorTReturn:
+    def __iter__(self) -> Iterator[int]:
+        ...
+
+
+class TypingIteratorReturn:
+    def __iter__(self) -> typing.Iterator:
+        ...
+
+
+class TypingIteratorTReturn:
+    def __iter__(self) -> typing.Iterator[int]:
+        ...
+
+
+class CollectionsIteratorReturn:
+    def __iter__(self) -> collections.abc.Iterator:
+        ...
+
+
+class CollectionsIteratorTReturn:
+    def __iter__(self) -> collections.abc.Iterator[int]:
+        ...
+
+
+class TypingAsyncIterableTReturn:
+    def __aiter__(self) -> typing.AsyncIterable[int]:
+        ...
+
+
+class TypingAsyncIterableReturn:
+    def __aiter__(self) -> typing.AsyncIterable:
+        ...

crates/ruff/resources/test/fixtures/flake8_pyi/PYI045.pyi

@@ -0,0 +1,49 @@
+import collections.abc
+import typing
+from collections.abc import Iterator, Iterable
+
+class NoReturn:
+    def __iter__(self): ...
+
+class TypingIterableTReturn:
+    def __iter__(self) -> typing.Iterable[int]: ...  # Error: PYI045
+    def not_iter(self) -> typing.Iterable[int]: ...
+
+class TypingIterableReturn:
+    def __iter__(self) -> typing.Iterable: ...  # Error: PYI045
+    def not_iter(self) -> typing.Iterable: ...
+
+class CollectionsIterableTReturn:
+    def __iter__(self) -> collections.abc.Iterable[int]: ...  # Error: PYI045
+    def not_iter(self) -> collections.abc.Iterable[int]: ...
+
+class CollectionsIterableReturn:
+    def __iter__(self) -> collections.abc.Iterable: ...  # Error: PYI045
+    def not_iter(self) -> collections.abc.Iterable: ...
+
+class IterableReturn:
+    def __iter__(self) -> Iterable: ...  # Error: PYI045
+
+class IteratorReturn:
+    def __iter__(self) -> Iterator: ...
+
+class IteratorTReturn:
+    def __iter__(self) -> Iterator[int]: ...
+
+class TypingIteratorReturn:
+    def __iter__(self) -> typing.Iterator: ...
+
+class TypingIteratorTReturn:
+    def __iter__(self) -> typing.Iterator[int]: ...
+
+class CollectionsIteratorReturn:
+    def __iter__(self) -> collections.abc.Iterator: ...
+
+class CollectionsIteratorTReturn:
+    def __iter__(self) -> collections.abc.Iterator[int]: ...
+
+class TypingAsyncIterableTReturn:
+    def __aiter__(self) -> typing.AsyncIterable[int]: ...  # Error: PYI045
+
+class TypingAsyncIterableReturn:
+    def __aiter__(self) -> typing.AsyncIterable: ...  # Error: PYI045

crates/ruff/src/checkers/ast/mod.rs

@@ -5289,7 +5289,8 @@ impl<'a> Checker<'a> {
             Rule::MissingReturnTypeClassMethod,
             Rule::AnyType,
         ]);
-        let enforce_stubs = self.is_stub && self.any_enabled(&[Rule::DocstringInStub]);
+        let enforce_stubs = self.is_stub
+            && self.any_enabled(&[Rule::DocstringInStub, Rule::IterMethodReturnIterable]);
         let enforce_docstrings = self.any_enabled(&[
             Rule::UndocumentedPublicModule,
             Rule::UndocumentedPublicClass,
@@ -5393,6 +5394,9 @@ impl<'a> Checker<'a> {
                     if self.enabled(Rule::DocstringInStub) {
                         flake8_pyi::rules::docstring_in_stubs(self, docstring);
                     }
+                    if self.enabled(Rule::IterMethodReturnIterable) {
+                        flake8_pyi::rules::iter_method_return_iterable(self, definition);
+                    }
                 }
             }
 

crates/ruff/src/codes.rs

@@ -594,6 +594,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Flake8Pyi, "033") => (RuleGroup::Unspecified, Rule::TypeCommentInStub),
         (Flake8Pyi, "042") => (RuleGroup::Unspecified, Rule::SnakeCaseTypeAlias),
         (Flake8Pyi, "043") => (RuleGroup::Unspecified, Rule::TSuffixedTypeAlias),
+        (Flake8Pyi, "045") => (RuleGroup::Unspecified, Rule::IterMethodReturnIterable),
         (Flake8Pyi, "048") => (RuleGroup::Unspecified, Rule::StubBodyMultipleStatements),
         (Flake8Pyi, "052") => (RuleGroup::Unspecified, Rule::UnannotatedAssignmentInStub),
 

crates/ruff/src/registry.rs

@@ -514,6 +514,7 @@ ruff_macros::register_rules!(
     rules::flake8_pyi::rules::AssignmentDefaultInStub,
     rules::flake8_pyi::rules::BadVersionInfoComparison,
     rules::flake8_pyi::rules::DocstringInStub,
+    rules::flake8_pyi::rules::IterMethodReturnIterable,
     rules::flake8_pyi::rules::DuplicateUnionMember,
     rules::flake8_pyi::rules::EllipsisInNonEmptyClassBody,
     rules::flake8_pyi::rules::NonEmptyStubBody,

crates/ruff/src/rules/flake8_pyi/mod.rs

@@ -26,6 +26,8 @@ mod tests {
     #[test_case(Rule::DuplicateUnionMember, Path::new("PYI016.pyi"))]
     #[test_case(Rule::EllipsisInNonEmptyClassBody, Path::new("PYI013.py"))]
     #[test_case(Rule::EllipsisInNonEmptyClassBody, Path::new("PYI013.pyi"))]
+    #[test_case(Rule::IterMethodReturnIterable, Path::new("PYI045.py"))]
+    #[test_case(Rule::IterMethodReturnIterable, Path::new("PYI045.pyi"))]
     #[test_case(Rule::NonEmptyStubBody, Path::new("PYI010.py"))]
     #[test_case(Rule::NonEmptyStubBody, Path::new("PYI010.pyi"))]
     #[test_case(Rule::PassInClassBody, Path::new("PYI012.py"))]

crates/ruff/src/rules/flake8_pyi/rules/iter_method_return_iterable.rs

@@ -0,0 +1,124 @@
+use rustpython_parser::ast;
+use rustpython_parser::ast::{Ranged, Stmt};
+
+use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::prelude::Expr;
+use ruff_python_semantic::definition::{Definition, Member, MemberKind};
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for `__iter__` methods in stubs that return `Iterable[T]` instead
+/// of an `Iterator[T]`.
+///
+/// ## Why is this bad?
+/// `__iter__` methods should always should return an `Iterator` of some kind,
+/// not an `Iterable`.
+///
+/// In Python, an `Iterator` is an object that has a `__next__` method, which
+/// provides a consistent interface for sequentially processing elements from
+/// a sequence or other iterable object. Meanwhile, an `Iterable` is an object
+/// with an `__iter__` method, which itself returns an `Iterator`.
+///
+/// Every `Iterator` is an `Iterable`, but not every `Iterable` is an `Iterator`.
+/// By returning an `Iterable` from `__iter__`, you may end up returning an
+/// object that doesn't implement `__next__`, which will cause a `TypeError`
+/// at runtime. For example, returning a `list` from `__iter__` will cause
+/// a `TypeError` when you call `__next__` on it, as a `list` is an `Iterable`,
+/// but not an `Iterator`.
+///
+/// ## Example
+/// ```python
+/// import collections.abc
+///
+///
+/// class Class:
+///     def __iter__(self) -> collections.abc.Iterable[str]:
+///         ...
+/// ```
+///
+/// Use instead:
+/// ```python
+/// import collections.abc
+///
+/// class Class:
+///     def __iter__(self) -> collections.abc.Iterator[str]:
+///         ...
+/// ```
+#[violation]
+pub struct IterMethodReturnIterable {
+    async_: bool,
+}
+
+impl Violation for IterMethodReturnIterable {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let IterMethodReturnIterable { async_ } = self;
+        if *async_ {
+            format!("`__aiter__` methods should return an `AsyncIterator`, not an `AsyncIterable`")
+        } else {
+            format!("`__iter__` methods should return an `Iterator`, not an `Iterable`")
+        }
+    }
+}
+
+/// PYI045
+pub(crate) fn iter_method_return_iterable(checker: &mut Checker, definition: &Definition) {
+    let Definition::Member(Member {
+        kind: MemberKind::Method,
+        stmt,
+        ..
+    }) = definition else {
+        return;
+    };
+
+    let Stmt::FunctionDef(ast::StmtFunctionDef {
+        name,
+        returns,
+        ..
+    }) = stmt else {
+        return;
+    };
+
+    let Some(returns) = returns else {
+        return;
+    };
+
+    let annotation = if let Expr::Subscript(ast::ExprSubscript { value, .. }) = returns.as_ref() {
+        // Ex) `Iterable[T]`
+        value
+    } else {
+        // Ex) `Iterable`, `typing.Iterable`
+        returns
+    };
+
+    let async_ = match name.as_str() {
+        "__iter__" => false,
+        "__aiter__" => true,
+        _ => return,
+    };
+
+    if checker
+        .semantic_model()
+        .resolve_call_path(annotation)
+        .map_or(false, |call_path| {
+            if async_ {
+                matches!(
+                    call_path.as_slice(),
+                    ["typing", "AsyncIterable"] | ["collections", "abc", "AsyncIterable"]
+                )
+            } else {
+                matches!(
+                    call_path.as_slice(),
+                    ["typing", "Iterable"] | ["collections", "abc", "Iterable"]
+                )
+            }
+        })
+    {
+        checker.diagnostics.push(Diagnostic::new(
+            IterMethodReturnIterable { async_ },
+            returns.range(),
+        ));
+    }
+}

crates/ruff/src/rules/flake8_pyi/rules/mod.rs

@@ -7,6 +7,9 @@ pub(crate) use duplicate_union_member::{duplicate_union_member, DuplicateUnionMe
 pub(crate) use ellipsis_in_non_empty_class_body::{
     ellipsis_in_non_empty_class_body, EllipsisInNonEmptyClassBody,
 };
+pub(crate) use iter_method_return_iterable::{
+    iter_method_return_iterable, IterMethodReturnIterable,
+};
 pub(crate) use non_empty_stub_body::{non_empty_stub_body, NonEmptyStubBody};
 pub(crate) use pass_in_class_body::{pass_in_class_body, PassInClassBody};
 pub(crate) use pass_statement_stub_body::{pass_statement_stub_body, PassStatementStubBody};
@@ -33,6 +36,7 @@ mod bad_version_info_comparison;
 mod docstring_in_stubs;
 mod duplicate_union_member;
 mod ellipsis_in_non_empty_class_body;
+mod iter_method_return_iterable;
 mod non_empty_stub_body;
 mod pass_in_class_body;
 mod pass_statement_stub_body;

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI045_PYI045.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+---
+

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI045_PYI045.pyi.snap

@@ -0,0 +1,61 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+---
+PYI045.pyi:9:27: PYI045 `__iter__` methods should return an `Iterator`, not an `Iterable`
+   |
+ 9 | class TypingIterableTReturn:
+10 |     def __iter__(self) -> typing.Iterable[int]: ...  # Error: PYI045
+   |                           ^^^^^^^^^^^^^^^^^^^^ PYI045
+11 |     def not_iter(self) -> typing.Iterable[int]: ...
+   |
+
+PYI045.pyi:13:27: PYI045 `__iter__` methods should return an `Iterator`, not an `Iterable`
+   |
+13 | class TypingIterableReturn:
+14 |     def __iter__(self) -> typing.Iterable: ...  # Error: PYI045
+   |                           ^^^^^^^^^^^^^^^ PYI045
+15 |     def not_iter(self) -> typing.Iterable: ...
+   |
+
+PYI045.pyi:17:27: PYI045 `__iter__` methods should return an `Iterator`, not an `Iterable`
+   |
+17 | class CollectionsIterableTReturn:
+18 |     def __iter__(self) -> collections.abc.Iterable[int]: ...  # Error: PYI045
+   |                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PYI045
+19 |     def not_iter(self) -> collections.abc.Iterable[int]: ...
+   |
+
+PYI045.pyi:21:27: PYI045 `__iter__` methods should return an `Iterator`, not an `Iterable`
+   |
+21 | class CollectionsIterableReturn:
+22 |     def __iter__(self) -> collections.abc.Iterable: ...  # Error: PYI045
+   |                           ^^^^^^^^^^^^^^^^^^^^^^^^ PYI045
+23 |     def not_iter(self) -> collections.abc.Iterable: ...
+   |
+
+PYI045.pyi:25:27: PYI045 `__iter__` methods should return an `Iterator`, not an `Iterable`
+   |
+25 | class IterableReturn:
+26 |     def __iter__(self) -> Iterable: ...  # Error: PYI045
+   |                           ^^^^^^^^ PYI045
+27 | 
+28 | class IteratorReturn:
+   |
+
+PYI045.pyi:46:28: PYI045 `__aiter__` methods should return an `AsyncIterator`, not an `AsyncIterable`
+   |
+46 | class TypingAsyncIterableTReturn:
+47 |     def __aiter__(self) -> typing.AsyncIterable[int]: ...  # Error: PYI045
+   |                            ^^^^^^^^^^^^^^^^^^^^^^^^^ PYI045
+48 | 
+49 | class TypingAsyncIterableReturn:
+   |
+
+PYI045.pyi:49:28: PYI045 `__aiter__` methods should return an `AsyncIterator`, not an `AsyncIterable`
+   |
+49 | class TypingAsyncIterableReturn:
+50 |     def __aiter__(self) -> typing.AsyncIterable: ...  # Error: PYI045
+   |                            ^^^^^^^^^^^^^^^^^^^^ PYI045
+   |
+
+