Skip to content

Commit

Permalink
fix: Fix building expressions (and string values) for yield and `yi…
Browse files Browse the repository at this point in the history
…eld from` statements
  • Loading branch information
pawamoy committed Feb 12, 2024
1 parent 4e3334e commit 439f65e
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 2 deletions.
9 changes: 7 additions & 2 deletions src/griffe/agents/nodes/_values.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,8 +291,12 @@ def _extract_usub(node: ast.USub, **kwargs: Any) -> str:

def _extract_yield(node: ast.Yield, **kwargs: Any) -> str:
if node.value is None:
return repr(None)
return _extract(node.value, **kwargs)
return "yield"
return f"yield {_extract(node.value, **kwargs)}"


def _extract_yield_from(node: ast.YieldFrom, **kwargs: Any) -> str:
return f"yield from {_extract(node.value, **kwargs)}"


_node_map: dict[type, Callable[[Any], str]] = {
Expand Down Expand Up @@ -353,6 +357,7 @@ def _extract_yield(node: ast.Yield, **kwargs: Any) -> str:
ast.UnaryOp: _extract_unaryop,
ast.USub: _extract_usub,
ast.Yield: _extract_yield,
ast.YieldFrom: _extract_yield_from,
}

# TODO: remove once Python 3.8 support is
Expand Down
17 changes: 17 additions & 0 deletions src/griffe/expressions.py
Original file line number Diff line number Diff line change
Expand Up @@ -724,6 +724,18 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]: # noqa: D102
yield from _yield(self.value, flat=flat)


@dataclass(eq=True, **dataclass_opts)
class ExprYieldFrom(Expr):
"""Yield statements like `yield from a`."""

value: str | Expr
"""Yielded-from value."""

def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
yield "yield from "
yield from _yield(self.value, flat=flat)


_unary_op_map = {
ast.Invert: "~",
ast.Not: "not ",
Expand Down Expand Up @@ -1002,6 +1014,10 @@ def _build_yield(node: ast.Yield, parent: Module | Class, **kwargs: Any) -> Expr
return ExprYield(None if node.value is None else _build(node.value, parent, **kwargs))


def _build_yield_from(node: ast.YieldFrom, parent: Module | Class, **kwargs: Any) -> Expr:
return ExprYieldFrom(_build(node.value, parent, **kwargs))


_node_map: dict[type, Callable[[Any, Module | Class], Expr]] = {
ast.Attribute: _build_attribute,
ast.BinOp: _build_binop,
Expand Down Expand Up @@ -1030,6 +1046,7 @@ def _build_yield(node: ast.Yield, parent: Module | Class, **kwargs: Any) -> Expr
ast.Tuple: _build_tuple,
ast.UnaryOp: _build_unaryop,
ast.Yield: _build_yield,
ast.YieldFrom: _build_yield_from,
}

# TODO: remove once Python 3.8 support is dropped
Expand Down
4 changes: 4 additions & 0 deletions tests/test_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,10 @@
"{a, b, c}",
"{a: b, c: d}",
"[a, b, c]",
# yields
"yield",
"yield a",
"yield from a",
]


Expand Down

0 comments on commit 439f65e

Please sign in to comment.