feat: Support unpacking typed dicts in signatures and docstrings

Issue-mkdocstrings-python-207: mkdocstrings/python#207
Issue-284: #284

Author: pawamoy

docs/extensions/built-in.md

@@ -5,3 +5,4 @@ Built-in extensions are maintained in Griffe's code base. They generally bring s
 Extension | Description
 --------- | -----------
 [`dataclasses`](built-in/dataclasses.md) | Support for [`dataclasses`][].
+[`unpack_typeddict`](built-in/unpack-typeddict.md) | Support for [`typing.Unpack`][] and [`typing.TypedDict`][].

docs/extensions/built-in/unpack-typeddict.md

@@ -0,0 +1,87 @@
+# `unpack`
+
+The `unpack_typeddict` extension adds support for [`Unpack`][typing.Unpack] and [`TypedDict`][typing.TypedDict] from the standard library. When enabled, it will add an `__init__` method to typed dictionaries, and expand `**kwargs: Unpack[...]` (the ellipsis being a [typed dict][typing.TypedDict] class) in function signatures to the relevant parameters, using the typed dict attributes. The extension will also update any Parameters section in the function docstring, to reflect the signature update.
+
+Example:
+
+```python
+from typing import TypedDict, Unpack
+
+
+class GreetKwargs(TypedDict):
+    name: str
+    """The name of a person to greet."""
+    shout: bool
+    """Whether to shout."""
+
+
+def greet(**kwargs: Unpack[GreetKwargs]) -> str:
+    """Greet someone.
+
+    Parameters:
+        **kwargs: Greet parameters.
+
+    Returns:
+        A message.
+    """
+    message = f"Hello {kwargs['name']}!"
+    if kwargs["shout"]:
+        return message.upper() + "!!"
+    return message
+```
+
+With the `unpack_typeddict` extension enabled, the data loaded by Griffe will be updated as follows:
+
+```python
+class GreetKwargs(TypedDict):
+    # Attributes removed from Griffe data (not from runtime class).
+
+    # Added by the `typeddict` extension to Griffe data (not to the runtime class):
+    def __init__(self, name: str, shout: bool) -> None:
+        """
+        Parameters:
+            name: The name of a person to greet.
+            shout: Whether to shout.
+        """
+
+
+def greet(*, name: str, shout: bool) -> str:
+    """Greet someone.
+
+    Parameters:
+        name: The name of a person to greet.
+        shout: Whether to shout.
+
+    Returns:
+        A message.
+    """
+```
+
+Thanks to this `__init__` method now appearing in the typed dictionary, tools like mkdocstrings can now render a proper signature for `GreetKwargs`.
+
+NOTE: Our example shows a Google-style docstring, but we actually insert a structured docstring section into the parsed data, which is style-agnostic, so it works with any docstring style.
+
+To enable the extension:
+
+=== "CLI"
+    ```console
+    $ griffe dump -e unpack_typeddict my_package
+    ```
+
+=== "Python"
+    ```python
+    import griffe
+
+    my_package = griffe.load("my_package", extensions=griffe.load_extensions("unpack_typeddict"))
+    ```
+
+=== "mkdocstrings"
+    ```yaml title="mkdocs.yml"
+    plugins:
+    - mkdocstrings:
+        handlers:
+          python:
+            options:
+              extensions:
+              - unpack_typeddict
+    ```

docs/reference/api/extensions.md

@@ -21,3 +21,7 @@
 ::: griffe.DataclassesExtension
     options:
         inherited_members: false
+
+::: griffe.UnpackTypedDictExtension
+    options:
+        inherited_members: false

duties.py

@@ -289,8 +289,12 @@ def check_api(ctx: Context, *cli_args: str) -> None:
             "griffe",
             search=["src"],
             color=True,
-            # YORE: Bump 2: Remove line.
-            extensions=["scripts/griffe_exts.py"],
+            extensions=[
+                "griffe_inherited_docstrings",
+                # YORE: Bump 2: Remove line.
+                "scripts/griffe_exts.py",
+                "unpack_typeddict",
+            ],
         ).add_args(*cli_args),
         title="Checking for API breaking changes",
         nofail=True,

mkdocs.yml

@@ -63,6 +63,7 @@ nav:
   - Built-in:
     - extensions/built-in.md
     - dataclasses: extensions/built-in/dataclasses.md
+    - unpack-typeddict: extensions/built-in/unpack-typeddict.md
   - Official:
     - extensions/official.md
     - autodocstringstyle: extensions/official/autodocstringstyle.md
@@ -232,6 +233,7 @@ plugins:
           - griffe_inherited_docstrings
           # YORE: Bump 2: Remove line.
           - scripts/griffe_exts.py
+          - unpack_typeddict
           heading_level: 2
           inherited_members: true
           merge_init_into_class: true

src/griffe/__init__.py

@@ -336,6 +336,7 @@
     load_extensions,
 )
 from griffe._internal.extensions.dataclasses import DataclassesExtension
+from griffe._internal.extensions.unpack_typeddict import UnpackTypedDictExtension
 from griffe._internal.finder import ModuleFinder, NamePartsAndPathType, NamePartsType, NamespacePackage, Package
 from griffe._internal.git import GitInfo, KnownGitService
 from griffe._internal.importer import dynamic_import, sys_path
@@ -558,6 +559,7 @@ def __getattr__(name: str) -> Any:
     "TypeParameters",
     "UnhandledEditableModuleError",
     "UnimportableModuleError",
+    "UnpackTypedDictExtension",
     "Visitor",
     "assert_git_repo",
     "ast_children",

src/griffe/_internal/extensions/base.py

@@ -520,6 +520,7 @@ def call(self, event: str, **kwargs: Any) -> None:
 
 builtin_extensions: set[str] = {
     "dataclasses",
+    "unpack_typeddict",
 }
 """The names of built-in Griffe extensions."""
 

src/griffe/_internal/extensions/unpack_typeddict.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+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
+
+
+def _update_docstring(func: Function, parameters: Iterable[Parameter], kwparam: Parameter | None = None) -> None:
+    if not func.docstring:
+        func.docstring = Docstring("", parent=func)
+    sections = func.docstring.parsed
+    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.
+        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 TypedDict attributes.
+    for param in parameters:
+        if param.name != "self":
+            params_section.value.append(
+                DocstringParameter(
+                    name=param.name,
+                    description=param.docstring.value if param.docstring else "",
+                    annotation=param.annotation,
+                    value=param.default,
+                ),
+            )
+
+
+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,
+            )
+            for attr in attrs
+        ),
+    )
+
+
+class UnpackTypedDictExtension(Extension):
+    """An extension to handle `Unpack[TypeDict]`."""
+
+    def on_class(self, *, cls: Class, **kwargs: Any) -> None:  # noqa: ARG002
+        """Add an `__init__` method to `TypedDict` classes if missing."""
+        for base in cls.bases:
+            if isinstance(base, Expr) and base.canonical_path in {"typing.TypedDict", "typing_extensions.TypedDict"}:
+                cls.labels.add("typed-dict")
+                break
+        else:
+            return
+
+        attributes = cls.attributes.values()
+
+        if "__init__" not in cls.members:
+            # Build the `__init__` method and add it to the class.
+            parameters = _params_from_attrs(attributes)
+            init = Function(name="__init__", parameters=parameters, returns="None")
+            cls.set_member("__init__", init)
+            # Update the `__init__` docstring.
+            _update_docstring(init, parameters)
+
+        # Remove attributes from the class, as they are now in the `__init__` method.
+        for attr in attributes:
+            cls.del_member(attr.name)
+
+    def on_function(self, *, func: Function, **kwargs: Any) -> None:  # noqa: ARG002
+        """Replace `**kwargs: Unpack[TypedDict]` parameters with the actual TypedDict attributes."""
+        # Find any `**kwargs: Unpack[TypedDict]` parameter.
+        for parameter in func.parameters:
+            if parameter.kind is ParameterKind.var_keyword:
+                annotation = parameter.annotation
+                if isinstance(annotation, ExprSubscript) and annotation.canonical_path in {
+                    "typing.Annotated",
+                    "typing_extensions.Annotated",
+                }:
+                    annotation = annotation.slice.elements[0]  # type: ignore[union-attr]
+                if isinstance(annotation, ExprSubscript) and annotation.canonical_path in {
+                    "typing.Unpack",
+                    "typing_extensions.Unpack",
+                }:
+                    slice_path = annotation.slice.canonical_path  # type: ignore[union-attr]
+                    typed_dict = func.modules_collection[slice_path]
+                    break
+        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())
+
+        # 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 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,
+                )

src/griffe/_internal/loader.py

@@ -190,7 +190,8 @@ def load(
         return self._post_load(top_module, obj_path)
 
     def _fire_load_events(self, obj: Object) -> None:
-        for member in obj.members.values():
+        # Wrapping in tuple() to avoid "dictionary changed size during iteration" errors.
+        for member in tuple(obj.members.values()):
             if member.is_alias:
                 self.extensions.call("on_alias", alias=member, loader=self)
                 continue

tests/test_api.py

@@ -18,8 +18,14 @@
 
 @pytest.fixture(name="loader", scope="module")
 def _fixture_loader() -> griffe.GriffeLoader:
-    # YORE: Bump 2: Regex-replace `extensions=.*` with `)` within line.
-    loader = griffe.GriffeLoader(extensions=griffe.load_extensions("scripts/griffe_exts.py"))
+    loader = griffe.GriffeLoader(
+        extensions=griffe.load_extensions(
+            "griffe_inherited_docstrings",
+            # YORE: Bump 2: Remove line.
+            "scripts/griffe_exts.py",
+            "unpack_typeddict",
+        ),
+    )
     loader.load("griffe")
     loader.resolve_aliases()
     return loader
@@ -150,7 +156,7 @@ def _public_path(obj: griffe.Object | griffe.Alias) -> bool:
 def test_api_matches_inventory(inventory: Inventory, public_objects: list[griffe.Object | griffe.Alias]) -> None:
     """All public objects are added to the inventory."""
     ignore_names = {"__getattr__", "__init__", "__repr__", "__str__", "__post_init__"}
-    ignore_paths = {"griffe.DataclassesExtension.*"}
+    ignore_paths = {"griffe.DataclassesExtension.*", "griffe.UnpackTypedDictExtension.*"}
     not_in_inventory = [
         f"{obj.relative_filepath}:{obj.lineno}: {obj.path}"
         for obj in public_objects