B901: Misses `yield` sub-expresisons

[MichaReiser]

A very contrived example but B901 misses yield expressions where they're not the expression of a ExprStmt.

def get_file_paths(file_types: Iterable[str] | None = None) -> Iterable[Path]:
    dir_path = Path(".")
    if file_types is None:
        return dir_path.glob("*")

    for file_type in file_types:
        (yield a for a in dir_path.glob(f"*.{file_type}"))

Note: we should make sure that we don't traverse into any lambda expressions...

[AlexWaygood]

def get_file_paths(file_types: Iterable[str] | None = None) -> Iterable[Path]:
    dir_path = Path(".")
    if file_types is None:
        return dir_path.glob("*")

    for file_type in file_types:
        (yield a for a in dir_path.glob(f"*.{file_type}"))

hmm, this isn't valid syntax:

>>> from typing import *
>>> from pathlib import Path
>>> def get_file_paths(file_types: Iterable[str] | None = None) -> Iterable[Path]:
...     dir_path = Path(".")
...     if file_types is None:
...         return dir_path.glob("*")
... 
...     for file_type in file_types:
...         (yield a for a in dir_path.glob(f"*.{file_type}"))
...         
  File "<python-input-2>", line 7
    (yield a for a in dir_path.glob(f"*.{file_type}"))
             ^^^
SyntaxError: invalid syntax

[MichaReiser]

hmm, I only tested with our parser. nested yield expressions are definitely valid :)

[MichaReiser]

Okay, found one

async def main():
    await (yield x)

[AlexWaygood]

Here's another function that is a generator function with a return statement and is not flagged by B901:

def f():
    x = yield
    print(x)
    return 42

This is a generator function that you can send values into:

>>> def f():
...     x = yield
...     print(x)
...     return 42
...     
>>> y = f()
>>> next(y)
>>> y.send('foo')
foo
Traceback (most recent call last):
  File "<python-input-13>", line 1, in <module>
    y.send('foo')
    ~~~~~~^^^^^^^
StopIteration: 42

Sending valued into a generator is, like returning values from a generator, quite an advanced feature. So if I saw a function like this in code review, I would assume the user knew what they were doing. That doesn't feel like a very principled reason not to emit B901 on it, though; maybe we should. Not sure.

[rabelmervin]

Hi @MichaReiser @AlexWaygood can I help you fix this issue ?

[MichaReiser]

Sure. You can find the code for this rule in this file

https://github.com/astral-sh/ruff/blob/b7e32b0a18487e8c9813e0d189d7a5d625ff5c3f/crates/ruff_linter/src/rules/flake8_bugbear/rules/return_in_generator.rs

[rabelmervin]

Thanks @MichaReiser excited to work on this !