fix: Correctly parenthesize expressions

Previously, stringifying expressions would concatenate the string representation of their parts without considering the order of operations, leading to incorrect removal of semantically necessary parentheses.

This commit changes the following things:

- the binding strength of operators (1 for walrus, 18 for atoms) were added
- parentheses are added when:
  - inner operation has a lower precedence than an outer one (e.g. `(a + b) * c`)
  - subexpression on the right of left associative operators have the same precedence (e.g. `(a - b) - c`)
  - subexpression on the left of right associative operators have the same precedence (e.g. `(a ** b) ** c`)
- special cases like power operator binding less when an arithmetic or bitwise operator on its right were implemented

PR-389: #389
Co-authored-by: Timothée Mazzucotelli <dev@pawamoy.fr>

Author: abc8747

src/_griffe/expressions.py

@@ -10,8 +10,8 @@
 import sys
 from dataclasses import dataclass
 from dataclasses import fields as getfields
+from enum import IntEnum, auto
 from functools import partial
-from itertools import zip_longest
 from typing import TYPE_CHECKING, Any, Callable
 
 from _griffe.agents.nodes.parameters import get_parameters
@@ -26,14 +26,82 @@
     from _griffe.models import Class, Function, Module
 
 
-def _yield(element: str | Expr | tuple[str | Expr, ...], *, flat: bool = True) -> Iterator[str | Expr]:
-    if isinstance(element, str):
-        yield element
+class _OperatorPrecedence(IntEnum):
+    # Adapted from:
+    #
+    # - https://docs.python.org/3/reference/expressions.html#operator-precedence
+    # - https://github.com/python/cpython/blob/main/Lib/_ast_unparse.py
+    # - https://github.com/astral-sh/ruff/blob/6abafcb56575454f2caeaa174efcb9fd0a8362b1/crates/ruff_python_ast/src/operator_precedence.rs
+
+    # The enum members are declared in ascending order of precedence.
+
+    # A virtual precedence level for contexts that provide their own grouping, like list brackets or
+    # function call parentheses. This ensures parentheses will never be added for the direct children of these nodes.
+    # NOTE: `ruff_python_formatter::expression::parentheses`'s state machine would be more robust
+    # but would introduce significant complexity.
+    NONE = auto()
+
+    YIELD = auto()  # `yield`, `yield from`
+    ASSIGN = auto()  # `target := expr`
+    STARRED = auto()  # `*expr` (omitted by Python docs, see ruff impl)
+    LAMBDA = auto()
+    IF_ELSE = auto()  # `expr if cond else expr`
+    OR = auto()
+    AND = auto()
+    NOT = auto()
+    COMP_MEMB_ID = auto()  # `<`, `<=`, `>`, `>=`, `!=`, `==`, `in`, `not in`, `is`, `is not`
+    BIT_OR = auto()  # `|`
+    BIT_XOR = auto()  # `^`
+    BIT_AND = auto()  # `&`
+    LEFT_RIGHT_SHIFT = auto()  # `<<`, `>>`
+    ADD_SUB = auto()  # `+`, `-`
+    MUL_DIV_REMAIN = auto()  # `*`, `@`, `/`, `//`, `%`
+    POS_NEG_BIT_NOT = auto()  # `+x`, `-x`, `~x`
+    EXPONENT = auto()  # `**`
+    AWAIT = auto()
+    CALL_ATTRIBUTE = auto()  # `x[index]`, `x[index:index]`, `x(arguments...)`, `x.attribute`
+    ATOMIC = auto()  # `(expressions...)`, `[expressions...]`, `{key: value...}`, `{expressions...}`
+
+
+def _yield(
+    element: str | Expr | tuple[str | Expr, ...],
+    *,
+    flat: bool = True,
+    is_left: bool = False,
+    outer_precedence: _OperatorPrecedence = _OperatorPrecedence.ATOMIC,
+) -> Iterator[str | Expr]:
+    if isinstance(element, Expr):
+        element_precedence = _get_precedence(element)
+        needs_parens = False
+        # Lower inner precedence, e.g. `(a + b) * c`, `+(10) < *(11)`.
+        if element_precedence < outer_precedence:
+            needs_parens = True
+        elif element_precedence == outer_precedence:
+            # Right-association, e.g. parenthesize left-hand side in `(a ** b) ** c`, (a if b else c) if d else e
+            is_right_assoc = isinstance(element, ExprIfExp) or (
+                isinstance(element, ExprBinOp) and element.operator == "**"
+            )
+            if is_right_assoc:
+                if is_left:
+                    needs_parens = True
+            # Left-association, e.g. parenthesize right-hand side in `a - (b - c)`.
+            elif isinstance(element, (ExprBinOp, ExprBoolOp)) and not is_left:
+                needs_parens = True
+
+        if needs_parens:
+            yield "("
+            if flat:
+                yield from element.iterate(flat=True)
+            else:
+                yield element
+            yield ")"
+        elif flat:
+            yield from element.iterate(flat=True)
+        else:
+            yield element
     elif isinstance(element, tuple):
         for elem in element:
-            yield from _yield(elem, flat=flat)
-    elif flat:
-        yield from element.iterate(flat=True)
+            yield from _yield(elem, flat=flat, outer_precedence=outer_precedence, is_left=is_left)
     else:
         yield element
 
@@ -44,14 +112,21 @@ def _join(
     *,
     flat: bool = True,
 ) -> Iterator[str | Expr]:
+    """Apply a separator between elements.
+
+    The caller is assumed to provide their own grouping
+    (e.g. lists, tuples, slice) and will prevent parentheses from being added.
+    """
     it = iter(elements)
     try:
-        yield from _yield(next(it), flat=flat)
+        # Since we are in a sequence, don't parenthesize items.
+        # Avoids [a + b, c + d] being serialized as [(a + b), (c + d)]
+        yield from _yield(next(it), flat=flat, outer_precedence=_OperatorPrecedence.NONE)
     except StopIteration:
         return
     for element in it:
-        yield from _yield(joint, flat=flat)
-        yield from _yield(element, flat=flat)
+        yield from _yield(joint, flat=flat, outer_precedence=_OperatorPrecedence.NONE)
+        yield from _yield(element, flat=flat, outer_precedence=_OperatorPrecedence.NONE)
 
 
 def _field_as_dict(
@@ -184,7 +259,11 @@ class ExprAttribute(Expr):
     """The different parts of the dotted chain."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield from _join(self.values, ".", flat=flat)
+        precedence = _get_precedence(self)
+        yield from _yield(self.values[0], flat=flat, outer_precedence=precedence, is_left=True)
+        for value in self.values[1:]:
+            yield "."
+            yield from _yield(value, flat=flat, outer_precedence=precedence)
 
     def append(self, value: ExprName) -> None:
         """Append a name to this attribute.
@@ -232,9 +311,14 @@ class ExprBinOp(Expr):
     """Right part."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield from _yield(self.left, flat=flat)
+        precedence = _get_precedence(self)
+        right_precedence = precedence
+        if self.operator == "**" and isinstance(self.right, ExprUnaryOp):
+            # Unary operators on the right have higher precedence, e.g. `a ** -b`.
+            right_precedence = _OperatorPrecedence(precedence - 1)
+        yield from _yield(self.left, flat=flat, outer_precedence=precedence, is_left=True)
         yield f" {self.operator} "
-        yield from _yield(self.right, flat=flat)
+        yield from _yield(self.right, flat=flat, outer_precedence=right_precedence, is_left=False)
 
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
@@ -248,7 +332,12 @@ class ExprBoolOp(Expr):
     """Operands."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield from _join(self.values, f" {self.operator} ", flat=flat)
+        precedence = _get_precedence(self)
+        it = iter(self.values)
+        yield from _yield(next(it), flat=flat, outer_precedence=precedence, is_left=True)
+        for value in it:
+            yield f" {self.operator} "
+            yield from _yield(value, flat=flat, outer_precedence=precedence, is_left=False)
 
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
@@ -267,7 +356,7 @@ def canonical_path(self) -> str:
         return self.function.canonical_path
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield from _yield(self.function, flat=flat)
+        yield from _yield(self.function, flat=flat, outer_precedence=_get_precedence(self))
         yield "("
         yield from _join(self.arguments, ", ", flat=flat)
         yield ")"
@@ -286,9 +375,11 @@ class ExprCompare(Expr):
     """Things compared."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield from _yield(self.left, flat=flat)
-        yield " "
-        yield from _join(zip_longest(self.operators, [], self.comparators, fillvalue=" "), " ", flat=flat)
+        precedence = _get_precedence(self)
+        yield from _yield(self.left, flat=flat, outer_precedence=precedence, is_left=True)
+        for op, comp in zip(self.operators, self.comparators):
+            yield f" {op} "
+            yield from _yield(comp, flat=flat, outer_precedence=precedence)
 
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
@@ -396,7 +487,8 @@ class ExprFormatted(Expr):
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
         yield "{"
-        yield from _yield(self.value, flat=flat)
+        # Prevent parentheses from being added, avoiding `{(1 + 1)}`
+        yield from _yield(self.value, flat=flat, outer_precedence=_OperatorPrecedence.NONE)
         yield "}"
 
 
@@ -429,11 +521,21 @@ class ExprIfExp(Expr):
     """Other expression."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield from _yield(self.body, flat=flat)
+        precedence = _get_precedence(self)
+        yield from _yield(self.body, flat=flat, outer_precedence=precedence, is_left=True)
         yield " if "
-        yield from _yield(self.test, flat=flat)
+        # If the test itself is another if/else, its precedence is the same, which will not give
+        # a parenthesis: force it.
+        test_outer_precedence = _OperatorPrecedence(precedence + 1)
+        yield from _yield(self.test, flat=flat, outer_precedence=test_outer_precedence)
         yield " else "
-        yield from _yield(self.orelse, flat=flat)
+        # If/else is right associative. For example, a nested if/else
+        # `a if b else c if d else e` is effectively `a if b else (c if d else e)`, so produce a
+        # flattened version without parentheses.
+        if isinstance(self.orelse, ExprIfExp):
+            yield from self.orelse.iterate(flat=flat)
+        else:
+            yield from _yield(self.orelse, flat=flat, outer_precedence=precedence, is_left=False)
 
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
@@ -557,7 +659,8 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
             if index < length:
                 yield ", "
         yield ": "
-        yield from _yield(self.body, flat=flat)
+        # Body of lambda should not have parentheses, avoiding `lambda: a.b`
+        yield from _yield(self.body, flat=flat, outer_precedence=_OperatorPrecedence.NONE)
 
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
@@ -685,11 +788,9 @@ class ExprNamedExpr(Expr):
     """Value."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield "("
         yield from _yield(self.target, flat=flat)
         yield " := "
         yield from _yield(self.value, flat=flat)
-        yield ")"
 
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
@@ -773,9 +874,10 @@ class ExprSubscript(Expr):
     """Slice part."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
-        yield from _yield(self.left, flat=flat)
+        yield from _yield(self.left, flat=flat, outer_precedence=_get_precedence(self))
         yield "["
-        yield from _yield(self.slice, flat=flat)
+        # Prevent parentheses from being added, avoiding `a[(b)]`
+        yield from _yield(self.slice, flat=flat, outer_precedence=_OperatorPrecedence.NONE)
         yield "]"
 
     @property
@@ -825,7 +927,9 @@ class ExprUnaryOp(Expr):
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
         yield self.operator
-        yield from _yield(self.value, flat=flat)
+        if self.operator == "not":
+            yield " "
+        yield from _yield(self.value, flat=flat, outer_precedence=_get_precedence(self))
 
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
@@ -858,7 +962,7 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
 
 _unary_op_map = {
     ast.Invert: "~",
-    ast.Not: "not ",
+    ast.Not: "not",
     ast.UAdd: "+",
     ast.USub: "-",
 }
@@ -897,6 +1001,65 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
     ast.NotIn: "not in",
 }
 
+# TODO: Support `ast.Await`.
+_precedence_map = {
+    # Literals and names.
+    ExprName: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprConstant: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprJoinedStr: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprFormatted: lambda _: _OperatorPrecedence.ATOMIC,
+    # Container displays.
+    ExprList: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprTuple: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprSet: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprDict: lambda _: _OperatorPrecedence.ATOMIC,
+    # Comprehensions are self-contained units that produce a container.
+    ExprListComp: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprSetComp: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprDictComp: lambda _: _OperatorPrecedence.ATOMIC,
+    ExprAttribute: lambda _: _OperatorPrecedence.CALL_ATTRIBUTE,
+    ExprSubscript: lambda _: _OperatorPrecedence.CALL_ATTRIBUTE,
+    ExprCall: lambda _: _OperatorPrecedence.CALL_ATTRIBUTE,
+    ExprUnaryOp: lambda e: {"not": _OperatorPrecedence.NOT}.get(e.operator, _OperatorPrecedence.POS_NEG_BIT_NOT),
+    ExprBinOp: lambda e: {
+        "**": _OperatorPrecedence.EXPONENT,
+        "*": _OperatorPrecedence.MUL_DIV_REMAIN,
+        "@": _OperatorPrecedence.MUL_DIV_REMAIN,
+        "/": _OperatorPrecedence.MUL_DIV_REMAIN,
+        "//": _OperatorPrecedence.MUL_DIV_REMAIN,
+        "%": _OperatorPrecedence.MUL_DIV_REMAIN,
+        "+": _OperatorPrecedence.ADD_SUB,
+        "-": _OperatorPrecedence.ADD_SUB,
+        "<<": _OperatorPrecedence.LEFT_RIGHT_SHIFT,
+        ">>": _OperatorPrecedence.LEFT_RIGHT_SHIFT,
+        "&": _OperatorPrecedence.BIT_AND,
+        "^": _OperatorPrecedence.BIT_XOR,
+        "|": _OperatorPrecedence.BIT_OR,
+    }[e.operator],
+    ExprBoolOp: lambda e: {"and": _OperatorPrecedence.AND, "or": _OperatorPrecedence.OR}[e.operator],
+    ExprCompare: lambda _: _OperatorPrecedence.COMP_MEMB_ID,
+    ExprIfExp: lambda _: _OperatorPrecedence.IF_ELSE,
+    ExprNamedExpr: lambda _: _OperatorPrecedence.ASSIGN,
+    ExprLambda: lambda _: _OperatorPrecedence.LAMBDA,
+    # NOTE: Ruff categorizes as atomic, but `(a for a in b).c` implies its less than `CALL_ATTRIBUTE`.
+    ExprGeneratorExp: lambda _: _OperatorPrecedence.LAMBDA,
+    ExprVarPositional: lambda _: _OperatorPrecedence.STARRED,
+    ExprVarKeyword: lambda _: _OperatorPrecedence.STARRED,
+    ExprYield: lambda _: _OperatorPrecedence.YIELD,
+    ExprYieldFrom: lambda _: _OperatorPrecedence.YIELD,
+    # These are not standalone, they appear in specific contexts where precendence is not a concern.
+    # NOTE: `for ... in ... if` part, not the whole `[...]`.
+    ExprComprehension: lambda _: _OperatorPrecedence.NONE,
+    ExprExtSlice: lambda _: _OperatorPrecedence.NONE,
+    ExprKeyword: lambda _: _OperatorPrecedence.NONE,
+    ExprParameter: lambda _: _OperatorPrecedence.NONE,
+    ExprSlice: lambda _: _OperatorPrecedence.NONE,
+}
+
+
+def _get_precedence(expr: Expr) -> _OperatorPrecedence:
+    return _precedence_map.get(type(expr), lambda _: _OperatorPrecedence.NONE)(expr)
+
 
 def _build_attribute(node: ast.Attribute, parent: Module | Class, **kwargs: Any) -> Expr:
     left = _build(node.value, parent, **kwargs)
@@ -1113,7 +1276,7 @@ def _build_subscript(
             "typing_extensions.Literal",
         }:
             literal_strings = True
-        slice = _build(
+        slice_expr = _build(
             node.slice,
             parent,
             parse_strings=True,
@@ -1122,8 +1285,8 @@ def _build_subscript(
             **kwargs,
         )
     else:
-        slice = _build(node.slice, parent, in_subscript=True, **kwargs)
-    return ExprSubscript(left, slice)
+        slice_expr = _build(node.slice, parent, in_subscript=True, **kwargs)
+    return ExprSubscript(left, slice_expr)
 
 
 def _build_tuple(