feat: Support expression modernization

Author: pawamoy

src/griffe/_internal/expressions.py

@@ -150,6 +150,13 @@ def _expr_as_dict(expression: Expr, **kwargs: Any) -> dict[str, Any]:
     return fields
 
 
+_modern_types = {
+    "typing.Tuple": "tuple",
+    "typing.Dict": "dict",
+    "typing.List": "list",
+    "typing.Set": "set",
+}
+
 # YORE: EOL 3.9: Remove block.
 _dataclass_opts: dict[str, bool] = {}
 if sys.version_info >= (3, 10):
@@ -265,6 +272,11 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
             yield "."
             yield from _yield(value, flat=flat, outer_precedence=precedence)
 
+    def modernize(self) -> ExprName | ExprAttribute:
+        if modern := _modern_types.get(self.canonical_path):
+            return ExprName(modern, parent=self.last.parent)
+        return self
+
     def append(self, value: ExprName) -> None:
         """Append a name to this attribute.
 
@@ -716,6 +728,11 @@ def __eq__(self, other: object) -> bool:
     def iterate(self, *, flat: bool = True) -> Iterator[ExprName]:  # noqa: ARG002
         yield self
 
+    def modernize(self) -> ExprName:
+        if modern := _modern_types.get(self.canonical_path):
+            return ExprName(modern, parent=self.parent)
+        return self
+
     @property
     def path(self) -> str:
         """The full, resolved name.
@@ -878,7 +895,7 @@ class ExprSubscript(Expr):
 
     left: str | Expr
     """Left part."""
-    slice: Expr
+    slice: str | Expr
     """Slice part."""
 
     def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
@@ -888,6 +905,33 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
         yield from _yield(self.slice, flat=flat, outer_precedence=_OperatorPrecedence.NONE)
         yield "]"
 
+    @staticmethod
+    def _to_binop(elements: Sequence[Expr], op: str) -> ExprBinOp:
+        if len(elements) == 2:  # noqa: PLR2004
+            left, right = elements
+            if isinstance(left, Expr):
+                left = left.modernize()
+            if isinstance(right, Expr):
+                right = right.modernize()
+            return ExprBinOp(left=left, operator=op, right=right)
+
+        left = ExprSubscript._to_binop(elements[:-1], op=op)
+        right = elements[-1]
+        if isinstance(right, Expr):
+            right = right.modernize()
+        return ExprBinOp(left=left, operator=op, right=right)
+
+    def modernize(self) -> ExprBinOp | ExprSubscript:
+        if self.canonical_path == "typing.Union":
+            return self._to_binop(self.slice.elements, op="|")  # type: ignore[union-attr]
+        if self.canonical_path == "typing.Optional":
+            left = self.slice if isinstance(self.slice, str) else self.slice.modernize()
+            return ExprBinOp(left=left, operator="|", right="None")
+        return ExprSubscript(
+            left=self.left if isinstance(self.left, str) else self.left.modernize(),
+            slice=self.slice if isinstance(self.slice, str) else self.slice.modernize(),
+        )
+
     @property
     def path(self) -> str:
         """The path of this subscript's left part."""
@@ -922,6 +966,12 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
         if not self.implicit:
             yield ")"
 
+    def modernize(self) -> ExprTuple:
+        return ExprTuple(
+            elements=[el if isinstance(el, str) else el.modernize() for el in self.elements],
+            implicit=self.implicit,
+        )
+
 
 # YORE: EOL 3.9: Replace `**_dataclass_opts` with `slots=True` within line.
 @dataclass(eq=True, **_dataclass_opts)

tests/test_expressions.py

@@ -95,6 +95,74 @@ def test_length_one_tuple_as_string() -> None:
         assert str(module["x"].value) == "('a',)"
 
 
+@pytest.mark.parametrize(
+    ("annotation", "modernized"),
+    [
+        ("Union[str, int, float]", "str | int | float"),
+        ("typing.Union[str, int, float]", "str | int | float"),
+        ("Union[Tuple[str, ...], Dict[str, int]]", "tuple[str, ...] | dict[str, int]"),
+        ("typing.Union[typing.Tuple[str, ...], typing.Dict[str, int]]", "tuple[str, ...] | dict[str, int]"),
+        ("Tuple[List[Dict[str, Set[str]]]]", "tuple[list[dict[str, set[str]]]]"),
+        ("typing.Tuple[typing.List[typing.Dict[str, typing.Set[str]]]]", "tuple[list[dict[str, set[str]]]]"),
+        ("Optional[Tuple[List[bool]]]", "tuple[list[bool]] | None"),
+        ("typing.Optional[typing.Tuple[typing.List[bool]]]", "tuple[list[bool]] | None"),
+    ],
+)
+def test_modernizing_specific_expressions(annotation: str, modernized: str) -> None:
+    """Modernize expressions correctly.
+
+    Parameters:
+        annotation: Original annotation (parametrized).
+        modernized: Expected modernized annotation (parametrized).
+    """
+    with temporary_visited_module(
+        f"""
+        import typing
+        from typing import Union, Optional, Tuple, Dict, List, Set, Literal
+        a: {annotation}
+        """,
+    ) as module:
+        expression = module["a"].annotation
+        assert str(expression.modernize()) == modernized
+
+
+@pytest.mark.parametrize(
+    "annotation",
+    [
+        "typing.Literal['s']",
+        "Literal['s']",
+    ],
+)
+def test_handling_modernization_without_crashing(annotation: str) -> None:
+    """Modernizing expressions never crashes.
+
+    Parameters:
+        annotation: Original annotation (parametrized).
+    """
+    with temporary_visited_module(
+        f"""
+        import typing
+        from typing import Union, Optional, Tuple, Dict, List, Set, Literal
+        a: {annotation}
+        """,
+    ) as module:
+        module["a"].annotation.modernize()
+
+
+@pytest.mark.parametrize("code", syntax_examples)
+def test_modernizing_idempotence(code: str) -> None:
+    """Modernize expressions that can't be modernized.
+
+    Parameters:
+        code: An expression (parametrized).
+    """
+    top_node = compile(code, filename="<>", mode="exec", flags=ast.PyCF_ONLY_AST, optimize=2)
+    expression = get_expression(top_node.body[0].value, parent=Module("module"))  # type: ignore[attr-defined]
+    modernized = expression.modernize()  # type: ignore[union-attr]
+    assert expression == modernized
+    assert str(expression) == str(modernized)
+
+
 def test_resolving_init_parameter() -> None:
     """Instance attribute values should resolve to matching parameters.