From 251adc20bf195d5366a2913e7ee9273d1f891a86 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20Mazzucotelli?= Date: Sat, 20 Sep 2025 14:36:19 +0200 Subject: [PATCH 1/2] refactor: Improve handling of typed dicts --- .../_internal/extensions/unpack_typeddict.py | 232 ++++++++++++++---- .../test_extensions/test_unpack_typeddict.py | 180 +++++++++++++- 2 files changed, 358 insertions(+), 54 deletions(-) diff --git a/src/griffe/_internal/extensions/unpack_typeddict.py b/src/griffe/_internal/extensions/unpack_typeddict.py index b25e5fa43..37cc6dee0 100644 --- a/src/griffe/_internal/extensions/unpack_typeddict.py +++ b/src/griffe/_internal/extensions/unpack_typeddict.py @@ -1,58 +1,188 @@ +# TODO: Support `extra_items=type`. +# TODO: Support `closed=True/False`. + from __future__ import annotations -from typing import TYPE_CHECKING, Any +import ast +from itertools import chain +from typing import TYPE_CHECKING, Any, TypedDict -from griffe._internal.docstrings.models import DocstringParameter, DocstringSectionParameters +from griffe._internal.docstrings.models import ( + DocstringParameter, + DocstringSectionParameters, +) from griffe._internal.enumerations import DocstringSectionKind, ParameterKind from griffe._internal.expressions import Expr, ExprSubscript from griffe._internal.extensions.base import Extension from griffe._internal.models import Class, Docstring, Function, Parameter, Parameters if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Iterable, Iterator + + +class _TypedDictAttr(TypedDict): + name: str + annotation: str | Expr | None + docstring: Docstring | None + + +def _unwrap_annotation(annotation: str | Expr | None, *, default_required: bool) -> tuple[str | Expr | None, bool]: + required = default_required + + # Annotations can be written ReadOnly[Required[T]] or Required[ReadOnly[T]], + # so we unwrap a first time here and a second time at the end. + if isinstance(annotation, ExprSubscript) and annotation.canonical_path in { + "typing.ReadOnly", + "typing_extensions.ReadOnly", + }: + annotation = annotation.slice # type: ignore[union-attr] + + # Unwrap `Required` and `NotRequired`, set `required` accordingly. + if isinstance(annotation, ExprSubscript): + if annotation.canonical_path in { + "typing.Required", + "typing_extensions.Required", + }: + annotation = annotation.slice # type: ignore[union-attr] + required = True + elif annotation.canonical_path in { + "typing.NotRequired", + "typing_extensions.NotRequired", + }: + annotation = annotation.slice # type: ignore[union-attr] + required = False + + # Unwrap `ReadOnly` a second time here. + if isinstance(annotation, ExprSubscript) and annotation.canonical_path in { + "typing.ReadOnly", + "typing_extensions.ReadOnly", + }: + annotation = annotation.slice # type: ignore[union-attr] + + return annotation, required + + +def _get_or_set_attrs(cls: Class) -> tuple[list[_TypedDictAttr], list[_TypedDictAttr]]: + if (attrs := cls.extra.get("unpack_typeddict", {}).get("_attributes")) is not None: + return attrs + + # Inspect `total` keyword argument to determine default requiredness. + default_required = True + for arg, value in cls.keywords.items(): + if arg == "total": + try: + total = ast.literal_eval(str(value)) + except (ValueError, SyntaxError): + break + if total is True: + default_required = True + elif total is False: + default_required = False + break + + # Extract attributes. + required_attrs = [] + optional_attrs = [] + for attr in cls.attributes.values(): + annotation, required = _unwrap_annotation(attr.annotation, default_required=default_required) + if required: + required_attrs.append( + _TypedDictAttr( + name=attr.name, + annotation=annotation, + docstring=attr.docstring, + ), + ) + else: + optional_attrs.append( + _TypedDictAttr( + name=attr.name, + annotation=annotation, + docstring=attr.docstring, + ), + ) + cls.extra["unpack_typeddict"]["_attributes"] = (required_attrs, optional_attrs) + return (required_attrs, optional_attrs) -def _update_docstring(func: Function, parameters: Iterable[Parameter], kwparam: Parameter | None = None) -> None: + +def _update_docstring( + func: Function, + required: Iterable[_TypedDictAttr], + optional: Iterable[_TypedDictAttr], + kwparam: Parameter | None = None, +) -> None: if not func.docstring: func.docstring = Docstring("", parent=func) + + params_section = None sections = func.docstring.parsed + + # Find existing "Parameters" section. section_gen = (section for section in sections if section.kind is DocstringSectionKind.parameters) - if kwparam and (params_section := next(section_gen, None)): - # Remove the `**kwargs` entry. + params_section = next(section_gen, None) + + # Pop original variadic keyword parameter from section. + if kwparam and params_section is not None: param_gen = (i for i, arg in enumerate(params_section.value) if arg.name.lstrip("*") == kwparam.name) if (kwarg_pos := next(param_gen, None)) is not None: params_section.value.pop(kwarg_pos) - else: - # Create a parameters section if none exists. - params_section = DocstringSectionParameters([]) - func.docstring.parsed.append(params_section) - # Add entries for all parameters. - for param in parameters: - if param.name != "self": + + # If we have required parameters, add them to the "Parameters" section. + if required: + # Create a "Parameters" section if none exists. + if params_section is None: + params_section = DocstringSectionParameters([]) + func.docstring.parsed.append(params_section) + + # Add required parameters to the section. + for attr in required: params_section.value.append( DocstringParameter( - name=param.name, - description=param.docstring.value if param.docstring else "", - annotation=param.annotation, - value=param.default, + name=attr["name"], + description=attr["docstring"].value if attr["docstring"] else "", + annotation=attr["annotation"], ), ) + # If we have optional parameters, add them to the "Parameters" section too, + # with a default value of `...`. + if optional: + # Create a "Parameters" section if none exists. + if params_section is None: + params_section = DocstringSectionParameters([]) + func.docstring.parsed.append(params_section) -def _params_from_attrs(attrs: Iterable[Any]) -> Parameters: - return Parameters( - Parameter(name="self", kind=ParameterKind.positional_or_keyword), - *( - Parameter( - name=attr.name, - annotation=attr.annotation, - kind=ParameterKind.keyword_only, - default=attr.value, - docstring=attr.docstring, + # Add optional parameters to the section. + for attr in optional: + params_section.value.append( + DocstringParameter( + name=attr["name"], + description=attr["docstring"].value if attr["docstring"] else "", + annotation=attr["annotation"], + value="...", + ), ) - for attr in attrs - ), - ) + + # TODO: Add `**kwargs` parameter if extra items are allowed. + + +def _params_from_attrs(required: Iterable[_TypedDictAttr], optional: Iterable[_TypedDictAttr]) -> Iterator[Parameter]: + for attr in required: + yield Parameter( + name=attr["name"], + annotation=attr["annotation"], + kind=ParameterKind.keyword_only, + docstring=attr["docstring"], + ) + for attr in optional: + yield Parameter( + name=attr["name"], + annotation=attr["annotation"], + kind=ParameterKind.keyword_only, + default="...", + docstring=attr["docstring"], + ) class UnpackTypedDictExtension(Extension): @@ -67,19 +197,23 @@ def on_class(self, *, cls: Class, **kwargs: Any) -> None: # noqa: ARG002 else: return - attributes = cls.attributes.values() + required, optional = _get_or_set_attrs(cls) if "__init__" not in cls.members: # Build the `__init__` method and add it to the class. - parameters = _params_from_attrs(attributes) + parameters = Parameters( + Parameter(name="self", kind=ParameterKind.positional_or_keyword), + *_params_from_attrs(required, optional), + ) + # TODO: Add `**kwargs` parameter if extra items are allowed. init = Function(name="__init__", parameters=parameters, returns="None") cls.set_member("__init__", init) # Update the `__init__` docstring. - _update_docstring(init, parameters) + _update_docstring(init, required, optional) # Remove attributes from the class, as they are now in the `__init__` method. - for attr in attributes: - cls.del_member(attr.name) + for attr in chain(required, optional): + cls.del_member(attr["name"]) def on_function(self, *, func: Function, **kwargs: Any) -> None: # noqa: ARG002 """Expand `**kwargs: Unpack[TypedDict]` in function signatures.""" @@ -102,26 +236,22 @@ def on_function(self, *, func: Function, **kwargs: Any) -> None: # noqa: ARG002 else: return - if "__init__" in typed_dict.members: - # The `__init__` was already generated: use its parameters. - parameters = typed_dict["__init__"].parameters - else: - # Fallback to building parameters from attributes. - parameters = _params_from_attrs(typed_dict.attributes.values()) + required, optional = _get_or_set_attrs(typed_dict) # Update any parameter section in the docstring. # We do this before updating the signature so that # parsing the docstring doesn't emit warnings. - _update_docstring(func, parameters, parameter) + _update_docstring(func, required, optional, parameter) # Update the function parameters. del func.parameters[parameter.name] - for param in parameters: - if param.name != "self": - func.parameters[param.name] = Parameter( - name=param.name, - annotation=param.annotation, - kind=ParameterKind.keyword_only, - default=param.default, - docstring=param.docstring, - ) + for param in _params_from_attrs(required, optional): + func.parameters[param.name] = Parameter( + name=param.name, + annotation=param.annotation, + kind=ParameterKind.keyword_only, + default=param.default, + docstring=param.docstring, + ) + + # TODO: Add `**kwargs` parameter if extra items are allowed. diff --git a/tests/test_extensions/test_unpack_typeddict.py b/tests/test_extensions/test_unpack_typeddict.py index 99c79450b..4aca4e357 100644 --- a/tests/test_extensions/test_unpack_typeddict.py +++ b/tests/test_extensions/test_unpack_typeddict.py @@ -18,8 +18,11 @@ class Kwargs(TypedDict): "pkg", {"__init__.py": code}, extensions=load_extensions("unpack_typeddict"), + docstring_parser="google", ) as pkg: td = pkg["Kwargs"] + + # Signature of the `__init__` method. assert "__init__" in td.members init = td["__init__"] assert len(init.parameters) == 3 @@ -32,6 +35,20 @@ class Kwargs(TypedDict): assert init.parameters["b"].docstring.value == "Docstring for b." assert init.returns == "None" + # Docstring and its "Parameters" section. + assert init.docstring + sections = init.docstring.parsed + assert len(sections) == 1 + params_section = sections[0] + assert params_section.kind is DocstringSectionKind.parameters + assert len(params_section.value) == 2 + assert params_section.value[0].name == "a" + assert params_section.value[0].description == "Docstring for a." + assert str(params_section.value[0].annotation) == "int" + assert params_section.value[1].name == "b" + assert params_section.value[1].description == "Docstring for b." + assert str(params_section.value[1].annotation) == "str" + def test_unpack_support() -> None: """Test our `Unpack` support.""" @@ -60,6 +77,7 @@ def func(**kwargs: Unpack[Kwargs]) -> None: ) as pkg: func = pkg["func"] + # Signature of the `func` function. assert len(func.parameters) == 2 assert "a" in func.parameters assert "b" in func.parameters @@ -67,10 +85,13 @@ def func(**kwargs: Unpack[Kwargs]) -> None: assert str(func.parameters["b"].annotation) == "str" assert func.parameters["a"].docstring.value == "Docstring for a." assert func.parameters["b"].docstring.value == "Docstring for b." + + # Docstring and its "Parameters" section. assert func.docstring is not None - params_section = next( - section for section in func.docstring.parsed if section.kind is DocstringSectionKind.parameters - ) + sections = func.docstring.parsed + assert len(sections) == 2 + params_section = sections[1] + assert params_section.kind is DocstringSectionKind.parameters assert len(params_section.value) == 2 param_a = next(param for param in params_section.value if param.name == "a") param_b = next(param for param in params_section.value if param.name == "b") @@ -78,3 +99,156 @@ def func(**kwargs: Unpack[Kwargs]) -> None: assert str(param_b.annotation) == "str" assert param_a.description == "Docstring for a." assert param_b.description == "Docstring for b." + + +def test_non_total_typeddict() -> None: + """Test our `TypedDict` support with non-total `TypedDict`s.""" + code = """ + from typing import TypedDict + + class Kwargs(TypedDict, total=False): + a: int + '''Docstring for a.''' + b: str + '''Docstring for b.''' + """ + with temporary_visited_package( + "pkg", + {"__init__.py": code}, + extensions=load_extensions("unpack_typeddict"), + docstring_parser="google", + ) as pkg: + td = pkg["Kwargs"] + + # Docstring and its "Parameters" section. + init = td["__init__"] + assert init.docstring + sections = init.docstring.parsed + assert len(sections) == 1 + params_section = sections[0] + assert params_section.kind is DocstringSectionKind.parameters + assert len(params_section.value) == 2 + assert params_section.value[0].name == "a" + assert params_section.value[0].description == "Docstring for a." + assert params_section.value[0].value == "..." + assert str(params_section.value[0].annotation) == "int" + assert params_section.value[1].name == "b" + assert params_section.value[1].description == "Docstring for b." + assert params_section.value[1].value == "..." + assert str(params_section.value[1].annotation) == "str" + + +def test_non_total_unpack() -> None: + """Test unpacking non-total `TypedDict`s.""" + code = """ + from typing import TypedDict, Unpack + + class Kwargs(TypedDict, total=False): + a: int + '''Docstring for a.''' + b: str + '''Docstring for b.''' + + def func(**kwargs: Unpack[Kwargs]) -> None: + '''A function.''' + """ + with temporary_visited_package( + "pkg", + {"__init__.py": code}, + extensions=load_extensions("unpack_typeddict"), + docstring_parser="google", + ) as pkg: + td = pkg["Kwargs"] + + # Docstring and its "Parameters" section. + init = td["__init__"] + assert init.docstring + sections = init.docstring.parsed + assert len(sections) == 1 + params_section = sections[0] + assert params_section.kind is DocstringSectionKind.parameters + assert len(params_section.value) == 2 + assert params_section.value[0].name == "a" + assert params_section.value[0].description == "Docstring for a." + assert params_section.value[0].value == "..." + assert str(params_section.value[0].annotation) == "int" + assert params_section.value[1].name == "b" + assert params_section.value[1].description == "Docstring for b." + assert params_section.value[1].value == "..." + assert str(params_section.value[1].annotation) == "str" + + +def test_explicit_requiredness() -> None: + """Test our `TypedDict` support with explicit requiredness.""" + code = """ + from typing import TypedDict + from typing_extensions import Required, NotRequired + + class Kwargs(TypedDict): + a: NotRequired[int] + '''Docstring for a.''' + b: Required[str] + '''Docstring for b.''' + """ + with temporary_visited_package( + "pkg", + {"__init__.py": code}, + extensions=load_extensions("unpack_typeddict"), + docstring_parser="google", + ) as pkg: + td = pkg["Kwargs"] + + # Signature of the `__init__` method. + assert "__init__" in td.members + init = td["__init__"] + assert len(init.parameters) == 3 + assert "self" in init.parameters + assert "a" in init.parameters + assert "b" in init.parameters + assert str(init.parameters["a"].annotation) == "int" + assert str(init.parameters["b"].annotation) == "str" + assert init.parameters["a"].docstring.value == "Docstring for a." + assert init.parameters["b"].docstring.value == "Docstring for b." + assert init.parameters["a"].default == "..." + assert init.parameters["b"].default is None + assert [p.name for p in init.parameters] == ["self", "b", "a"] + + + +def test_readonly_fields() -> None: + """Test our `TypedDict` support with `ReadOnly` fields.""" + code = """ + from typing import TypedDict + from typing_extensions import ReadOnly, Required, NotRequired + + class Kwargs(TypedDict): + a: ReadOnly[int] + '''Docstring for a.''' + b: ReadOnly[Required[str]] + '''Docstring for b.''' + c: Required[ReadOnly[float]] + '''Docstring for c.''' + """ + with temporary_visited_package( + "pkg", + {"__init__.py": code}, + extensions=load_extensions("unpack_typeddict"), + docstring_parser="google", + ) as pkg: + td = pkg["Kwargs"] + + # Signature of the `__init__` method. + assert "__init__" in td.members + init = td["__init__"] + assert len(init.parameters) == 4 + assert "self" in init.parameters + assert "a" in init.parameters + assert "b" in init.parameters + assert "c" in init.parameters + assert str(init.parameters["a"].annotation) == "int" + assert str(init.parameters["b"].annotation) == "str" + assert str(init.parameters["c"].annotation) == "float" + assert init.parameters["a"].docstring.value == "Docstring for a." + assert init.parameters["b"].docstring.value == "Docstring for b." + assert init.parameters["c"].docstring.value == "Docstring for c." + assert init.returns == "None" From 93a3af01f80a8bbd1d8b496549bec1644b923e46 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20Mazzucotelli?= Date: Sat, 20 Sep 2025 14:42:41 +0200 Subject: [PATCH 2/2] fixup! refactor: Improve handling of typed dicts --- src/griffe/_internal/extensions/unpack_typeddict.py | 8 ++++---- tests/test_extensions/test_unpack_typeddict.py | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/griffe/_internal/extensions/unpack_typeddict.py b/src/griffe/_internal/extensions/unpack_typeddict.py index 37cc6dee0..162cbbf69 100644 --- a/src/griffe/_internal/extensions/unpack_typeddict.py +++ b/src/griffe/_internal/extensions/unpack_typeddict.py @@ -35,7 +35,7 @@ def _unwrap_annotation(annotation: str | Expr | None, *, default_required: bool) "typing.ReadOnly", "typing_extensions.ReadOnly", }: - annotation = annotation.slice # type: ignore[union-attr] + annotation = annotation.slice # Unwrap `Required` and `NotRequired`, set `required` accordingly. if isinstance(annotation, ExprSubscript): @@ -43,13 +43,13 @@ def _unwrap_annotation(annotation: str | Expr | None, *, default_required: bool) "typing.Required", "typing_extensions.Required", }: - annotation = annotation.slice # type: ignore[union-attr] + annotation = annotation.slice required = True elif annotation.canonical_path in { "typing.NotRequired", "typing_extensions.NotRequired", }: - annotation = annotation.slice # type: ignore[union-attr] + annotation = annotation.slice required = False # Unwrap `ReadOnly` a second time here. @@ -57,7 +57,7 @@ def _unwrap_annotation(annotation: str | Expr | None, *, default_required: bool) "typing.ReadOnly", "typing_extensions.ReadOnly", }: - annotation = annotation.slice # type: ignore[union-attr] + annotation = annotation.slice return annotation, required diff --git a/tests/test_extensions/test_unpack_typeddict.py b/tests/test_extensions/test_unpack_typeddict.py index 4aca4e357..fa24d9036 100644 --- a/tests/test_extensions/test_unpack_typeddict.py +++ b/tests/test_extensions/test_unpack_typeddict.py @@ -214,7 +214,6 @@ class Kwargs(TypedDict): assert [p.name for p in init.parameters] == ["self", "b", "a"] - def test_readonly_fields() -> None: """Test our `TypedDict` support with `ReadOnly` fields.""" code = """