feat: Parse Sphinx parameter types as expressions

PR-392: #392

Author: echoix

src/_griffe/docstrings/sphinx.py

@@ -22,7 +22,7 @@
     DocstringSectionReturns,
     DocstringSectionText,
 )
-from _griffe.docstrings.utils import docstring_warning
+from _griffe.docstrings.utils import docstring_warning, parse_docstring_annotation
 
 if TYPE_CHECKING:
     from _griffe.expressions import Expr
@@ -75,7 +75,7 @@ class _ParsedValues:
 
     description: list[str] = field(default_factory=list)
     parameters: dict[str, DocstringParameter] = field(default_factory=dict)
-    param_types: dict[str, str] = field(default_factory=dict)
+    param_types: dict[str, str | Expr] = field(default_factory=dict)
     attributes: dict[str, DocstringAttribute] = field(default_factory=dict)
     attribute_types: dict[str, str] = field(default_factory=dict)
     exceptions: list[DocstringRaise] = field(default_factory=list)
@@ -145,7 +145,10 @@ def _read_parameter(
         # no type info
         name = parsed_directive.directive_parts[1]
     elif len(parsed_directive.directive_parts) == 3:  # noqa: PLR2004
-        directive_type = parsed_directive.directive_parts[1]
+        directive_type = parse_docstring_annotation(
+            parsed_directive.directive_parts[1],
+            docstring,
+        )
         name = parsed_directive.directive_parts[2]
     else:
         if warnings:
@@ -191,7 +194,7 @@ def _determine_param_default(docstring: Docstring, name: str) -> str | None:
 def _determine_param_annotation(
     docstring: Docstring,
     name: str,
-    directive_type: str | None,
+    directive_type: str | Expr | None,
     parsed_values: _ParsedValues,
     *,
     warnings: bool = True,
@@ -234,7 +237,8 @@ def _read_parameter_type(
     parsed_directive = _parse_directive(docstring, offset, warnings=warnings)
     if parsed_directive.invalid:
         return parsed_directive.next_index
-    param_type = _consolidate_descriptive_type(parsed_directive.value.strip())
+    param_type_str = _consolidate_descriptive_type(parsed_directive.value.strip())
+    param_type = parse_docstring_annotation(param_type_str, docstring)
 
     if len(parsed_directive.directive_parts) == 2:  # noqa: PLR2004
         param_name = parsed_directive.directive_parts[1]

tests/test_docstrings/test_sphinx.py

@@ -15,6 +15,9 @@
     DocstringRaise,
     DocstringReturn,
     DocstringSectionKind,
+    Expr,
+    ExprBinOp,
+    ExprName,
     Function,
     Module,
     Parameter,
@@ -231,28 +234,62 @@ def test_parse__param_field_docs_type__param_section_with_type(parse_sphinx: Par
     assert actual.as_dict() == expected.as_dict()
 
 
-def test_parse__param_field_type_field__param_section_with_type(parse_sphinx: ParserType) -> None:
+@pytest.mark.parametrize("type_", ["str", "int"])
+def test_parse__param_field_type_field__param_section_with_type(parse_sphinx: ParserType, type_: str) -> None:
     """Parse parameters with separated types.
 
     Parameters:
         parse_sphinx: Fixture parser.
+        type_: The type to use in the type directive.
     """
     docstring = f"""
         Docstring with line continuation.
 
-        :param foo: {SOME_TEXT}
-        :type foo: str
+        :param {SOME_NAME}: {SOME_TEXT}
+        :type {SOME_NAME}: {type_}
     """
 
     sections, _ = parse_sphinx(docstring)
     assert len(sections) == 2
     assert sections[1].kind is DocstringSectionKind.parameters
     actual = sections[1].value[0]
-    expected = DocstringParameter(SOME_NAME, annotation="str", description=SOME_TEXT)
+    expected = DocstringParameter(SOME_NAME, annotation=f"{type_}", description=SOME_TEXT)
     assert isinstance(actual, type(expected))
     assert actual.as_dict() == expected.as_dict()
 
 
+@pytest.mark.parametrize("type_", ["str", "int"])
+def test_parse__param_field_type_field__param_section_with_type_with_parent(
+    parse_sphinx: ParserType,
+    type_: str,
+) -> None:
+    """Parse parameters with separated types.
+
+    Parameters:
+        parse_sphinx: Fixture parser.
+        type_: The type to use in the type directive.
+    """
+    docstring = f"""
+        Docstring with line continuation.
+
+        :param {SOME_NAME}: {SOME_TEXT}
+        :type {SOME_NAME}: {type_}
+    """
+    parent_fn = Function("func", parameters=Parameters(Parameter(SOME_NAME)))
+    sections, _ = parse_sphinx(docstring, parent=parent_fn)
+    assert len(sections) == 2
+    assert sections[1].kind is DocstringSectionKind.parameters
+    actual = sections[1].value[0]
+    expected_annotation = ExprName(name=f"{type_}")
+    expected = DocstringParameter(SOME_NAME, annotation=expected_annotation, description=SOME_TEXT)
+    assert isinstance(actual, type(expected))
+    assert actual.as_dict() == expected.as_dict()
+    assert isinstance(actual.annotation, type(expected.annotation))
+    assert isinstance(actual.annotation, ExprName)
+    assert isinstance(actual.annotation, Expr)
+    assert actual.annotation.as_dict() == expected_annotation.as_dict()
+
+
 def test_parse__param_field_type_field_first__param_section_with_type(parse_sphinx: ParserType) -> None:
     """Parse parameters with separated types first.
 
@@ -275,6 +312,33 @@ def test_parse__param_field_type_field_first__param_section_with_type(parse_sphi
     assert actual.as_dict() == expected.as_dict()
 
 
+def test_parse__param_field_type_field_first__param_section_with_type_with_parent(parse_sphinx: ParserType) -> None:
+    """Parse parameters with separated types first.
+
+    Parameters:
+        parse_sphinx: Fixture parser.
+    """
+    docstring = f"""
+        Docstring with line continuation.
+
+        :type {SOME_NAME}: str
+        :param {SOME_NAME}: {SOME_TEXT}
+    """
+    parent_fn = Function("func", parameters=Parameters(Parameter(SOME_NAME)))
+    sections, _ = parse_sphinx(docstring, parent=parent_fn)
+    assert len(sections) == 2
+    assert sections[1].kind is DocstringSectionKind.parameters
+    actual = sections[1].value[0]
+    expected_annotation = ExprName("str", parent=Class("C"))
+    expected = DocstringParameter(SOME_NAME, annotation=expected_annotation, description=SOME_TEXT)
+    assert isinstance(actual, type(expected))
+    assert actual.as_dict() == expected.as_dict()
+    assert isinstance(actual.annotation, type(expected.annotation))
+    assert isinstance(actual.annotation, ExprName)
+    assert isinstance(actual.annotation, Expr)
+    assert actual.annotation.as_dict() == expected_annotation.as_dict()
+
+
 @pytest.mark.parametrize("union", ["str or None", "None or str", "str or int", "str or int or float"])
 def test_parse__param_field_type_field_or_none__param_section_with_optional(
     parse_sphinx: ParserType,
@@ -302,6 +366,47 @@ def test_parse__param_field_type_field_or_none__param_section_with_optional(
     assert actual.as_dict() == expected.as_dict()
 
 
+@pytest.mark.parametrize(
+    ("union", "expected_annotation"),
+    [
+        ("str or None", ExprBinOp(ExprName("str"), "|", "None")),
+        ("None or str", ExprBinOp("None", "|", ExprName("str"))),
+        ("str or int", ExprBinOp(ExprName("str"), "|", ExprName("int"))),
+        ("str or int or float", ExprBinOp(ExprBinOp(ExprName("str"), "|", ExprName("int")), "|", ExprName("float"))),
+    ],
+)
+def test_parse__param_field_type_field_or_none__param_section_with_optional_with_parent(
+    parse_sphinx: ParserType,
+    union: str,
+    expected_annotation: Expr,
+) -> None:
+    """Parse parameters with separated union types.
+
+    Parameters:
+        parse_sphinx: Fixture parser.
+        union: A parametrized union type.
+        expected_annotation: The expected annotation as an expression
+    """
+    docstring = f"""
+        Docstring with line continuation.
+
+        :param {SOME_NAME}: {SOME_TEXT}
+        :type {SOME_NAME}: {union}
+    """
+
+    parent_fn = Function("func", parameters=Parameters(Parameter(SOME_NAME)))
+    sections, _ = parse_sphinx(docstring, parent=parent_fn)
+    assert len(sections) == 2
+    assert sections[1].kind is DocstringSectionKind.parameters
+    actual = sections[1].value[0]
+    expected = DocstringParameter(SOME_NAME, annotation=expected_annotation, description=SOME_TEXT)
+    assert isinstance(actual, type(expected))
+    assert actual.as_dict() == expected.as_dict()
+    assert isinstance(actual.annotation, type(expected.annotation))
+    assert isinstance(actual.annotation, Expr)
+    assert actual.annotation.as_dict() == expected_annotation.as_dict()
+
+
 def test_parse__param_field_annotate_type__param_section_with_type(parse_sphinx: ParserType) -> None:
     """Parse a simple docstring.