Introduce Y058: Use (Async)Iterator, not (Async)Generator, for si…

…mple `__(a)iter__` methods (#428)

Closes #423

CHANGELOG.md

@@ -6,6 +6,9 @@
   attributes has been removed. Use `flake8 --version` from the command line, or
   `importlib.metadata.version("flake8_pyi")` at runtime, to determine the
   version of `flake8-pyi` installed at runtime.
+* Introduce Y058: Use `Iterator` rather than `Generator` as the return value
+  for simple `__iter__` methods, and `AsyncIterator` rather than
+  `AsyncGenerator` as the return value for simple `__aiter__` methods.
 
 ## 23.10.0
 

ERRORCODES.md

@@ -60,6 +60,7 @@ The following warnings are currently emitted by default:
 | Y055 | Unions of the form `type[X] \| type[Y]` can be simplified to `type[X \| Y]`. Similarly, `Union[type[X], type[Y]]` can be simplified to `type[Union[X, Y]]`.
 | Y056 | Do not call methods such as `.append()`, `.extend()` or `.remove()` on `__all__`. Different type checkers have varying levels of support for calling these methods on `__all__`. Use `+=` instead, which is known to be supported by all major type checkers.
 | Y057 | Do not use `typing.ByteString` or `collections.abc.ByteString`. These types have unclear semantics, and are deprecated; use  `typing_extensions.Buffer` or a union such as `bytes \| bytearray \| memoryview` instead. See [PEP 688](https://peps.python.org/pep-0688/) for more details.
+| Y058 | Use `Iterator` rather than `Generator` as the return value for simple `__iter__` methods, and `AsyncIterator` rather than `AsyncGenerator` as the return value for simple `__aiter__` methods. Using `(Async)Iterator` for these methods is simpler and more elegant, and reflects the fact that the precise kind of iterator returned from an `__iter__` method is usually an implementation detail that could change at any time, and should not be relied upon.
 
 ## Warnings disabled by default
 

pyi.py

@@ -341,6 +341,12 @@ def _is_object(node: ast.AST | None, name: str, *, from_: Container[str]) -> boo
 _is_Protocol = partial(_is_object, name="Protocol", from_=_TYPING_MODULES)
 _is_NoReturn = partial(_is_object, name="NoReturn", from_=_TYPING_MODULES)
 _is_Final = partial(_is_object, name="Final", from_=_TYPING_MODULES)
+_is_Generator = partial(
+    _is_object, name="Generator", from_=_TYPING_MODULES | {"collections.abc"}
+)
+_is_AsyncGenerator = partial(
+    _is_object, name="AsyncGenerator", from_=_TYPING_MODULES | {"collections.abc"}
+)
 
 
 def _is_object_or_Unused(node: ast.expr | None) -> bool:
@@ -1700,6 +1706,15 @@ def _Y034_error(
         )
         self.error(node, error_message)
 
+    def _Y058_error(
+        self, node: ast.FunctionDef, args: list[ast.arg], example_returns: str
+    ) -> None:
+        assert node.name in {"__iter__", "__aiter__"}
+        good_cls = "Iterator" if node.name == "__iter__" else "AsyncIterator"
+        example = f"def {node.name}({args[0].arg}) -> {example_returns}: ..."
+        msg = Y058.format(iter_method=node.name, good_cls=good_cls, example=example)
+        self.error(node, msg)
+
     def _check_iter_returns(
         self, node: ast.FunctionDef, returns: ast.expr | None
     ) -> None:
@@ -1710,6 +1725,24 @@ def _check_iter_returns(
                 iter_method="__iter__", good_cls="Iterator", bad_cls="Iterable"
             )
             self.error(node, msg)
+            return
+        non_kw_only_args = node.args.posonlyargs + node.args.args
+        if len(non_kw_only_args) == 1 and not node.args.kwonlyargs:
+            if _is_Generator(returns):
+                self._Y058_error(node, non_kw_only_args, "Iterator")
+            elif (
+                isinstance(returns, ast.Subscript)
+                and _is_Generator(returns.value)
+                and isinstance(returns.slice, ast.Tuple)
+            ):
+                elts = returns.slice.elts
+                if (
+                    len(elts) == 3
+                    and (_is_Any(elts[1]) or _is_None(elts[1]))
+                    and (_is_Any(elts[2]) or _is_None(elts[2]))
+                ):
+                    example_returns = f"Iterator[{unparse(returns.slice.elts[0])}]"
+                    self._Y058_error(node, non_kw_only_args, example_returns)
 
     def _check_aiter_returns(
         self, node: ast.FunctionDef, returns: ast.expr | None
@@ -1723,6 +1756,20 @@ def _check_aiter_returns(
                 bad_cls="AsyncIterable",
             )
             self.error(node, msg)
+            return
+        non_kw_only_args = node.args.posonlyargs + node.args.args
+        if len(non_kw_only_args) == 1 and not node.args.kwonlyargs:
+            if _is_AsyncGenerator(returns):
+                self._Y058_error(node, non_kw_only_args, "AsyncIterator")
+            elif (
+                isinstance(returns, ast.Subscript)
+                and _is_AsyncGenerator(returns.value)
+                and isinstance(returns.slice, ast.Tuple)
+            ):
+                elts = returns.slice.elts
+                if len(elts) == 2 and (_is_Any(elts[1]) or _is_None(elts[1])):
+                    example_returns = f"AsyncIterator[{unparse(returns.slice.elts[0])}]"
+                    self._Y058_error(node, non_kw_only_args, example_returns)
 
     def _visit_synchronous_method(self, node: ast.FunctionDef) -> None:
         method_name = node.name
@@ -2125,6 +2172,10 @@ def parse_options(options: argparse.Namespace) -> None:
 Y057 = (
     "Y057 Do not use {module}.ByteString, which has unclear semantics and is deprecated"
 )
+Y058 = (
+    'Y058 Use "{good_cls}" as the return value for simple "{iter_method}" methods, '
+    'e.g. "{example}"'
+)
 Y090 = (
     'Y090 "{original}" means '
     '"a tuple of length 1, in which the sole element is of type {typ!r}". '

tests/classdefs.pyi

@@ -5,7 +5,14 @@ import builtins
 import collections.abc
 import typing
 from abc import abstractmethod
-from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator
+from collections.abc import (
+    AsyncGenerator,
+    AsyncIterable,
+    AsyncIterator,
+    Generator,
+    Iterable,
+    Iterator,
+)
 from typing import Any, overload
 
 import typing_extensions
@@ -94,12 +101,30 @@ class BadIterator4(Iterator[int]):
 class IteratorReturningIterable:
     def __iter__(self) -> Iterable[str]: ...  # Y045 "__iter__" methods should return an Iterator, not an Iterable
 
+class IteratorReturningSimpleGenerator1:
+    def __iter__(self) -> Generator: ...  # Y058 Use "Iterator" as the return value for simple "__iter__" methods, e.g. "def __iter__(self) -> Iterator: ..."
+
+class IteratorReturningSimpleGenerator2:
+    def __iter__(self) -> collections.abc.Generator[str, Any, None]: ...  # Y058 Use "Iterator" as the return value for simple "__iter__" methods, e.g. "def __iter__(self) -> Iterator[str]: ..."
+
+class IteratorReturningComplexGenerator:
+    def __iter__(self) -> Generator[str, int, bytes]: ...
+
 class BadAsyncIterator(collections.abc.AsyncIterator[str]):
     def __aiter__(self) -> typing.AsyncIterator[str]: ...  # Y034 "__aiter__" methods in classes like "BadAsyncIterator" usually return "self" at runtime. Consider using "typing_extensions.Self" in "BadAsyncIterator.__aiter__", e.g. "def __aiter__(self) -> Self: ..."  # Y022 Use "collections.abc.AsyncIterator[T]" instead of "typing.AsyncIterator[T]" (PEP 585 syntax)
 
 class AsyncIteratorReturningAsyncIterable:
     def __aiter__(self) -> AsyncIterable[str]: ...  # Y045 "__aiter__" methods should return an AsyncIterator, not an AsyncIterable
 
+class AsyncIteratorReturningSimpleAsyncGenerator1:
+    def __aiter__(self) -> AsyncGenerator: ...  # Y058 Use "AsyncIterator" as the return value for simple "__aiter__" methods, e.g. "def __aiter__(self) -> AsyncIterator: ..."
+
+class AsyncIteratorReturningSimpleAsyncGenerator2:
+    def __aiter__(self) -> collections.abc.AsyncGenerator[str, Any]: ...  # Y058 Use "AsyncIterator" as the return value for simple "__aiter__" methods, e.g. "def __aiter__(self) -> AsyncIterator[str]: ..."
+
+class AsyncIteratorReturningComplexAsyncGenerator:
+    def __aiter__(self) -> AsyncGenerator[str, int]: ...
+
 class Abstract(Iterator[str]):
     @abstractmethod
     def __iter__(self) -> Iterator[str]: ...