wip

Author: pawamoy

src/griffe/_internal/extensions/unpack_typeddict.py

@@ -1,15 +1,14 @@
-# TODO: Support `typing.ReadOnly`.
 # TODO: Support `extra_items=type`.
 # TODO: Support `closed=True/False`.
 
 from __future__ import annotations
 
 import ast
+from itertools import chain
 from typing import TYPE_CHECKING, Any, TypedDict
 
 from griffe._internal.docstrings.models import (
     DocstringParameter,
-    DocstringSectionOtherParameters,
     DocstringSectionParameters,
 )
 from griffe._internal.enumerations import DocstringSectionKind, ParameterKind
@@ -18,21 +17,56 @@
 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
-    required: bool
 
 
-def _get_or_set_attrs(cls: Class) -> list[_TypedDictAttr]:
+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
 
-    attrs = []
+    # Inspect `total` keyword argument to determine default requiredness.
     default_required = True
     for arg, value in cls.keywords.items():
         if arg == "total":
@@ -45,61 +79,54 @@ def _get_or_set_attrs(cls: Class) -> list[_TypedDictAttr]:
             elif total is False:
                 default_required = False
             break
+
+    # Extract attributes.
+    required_attrs = []
+    optional_attrs = []
     for attr in cls.attributes.values():
-        annotation = attr.annotation
-        required = default_required
-        if isinstance(annotation, ExprSubscript):
-            if annotation.canonical_path in {
-                "typing.Required",
-                "typing_extensions.Required",
-            }:
-                annotation = annotation.slice.elements[0]  # type: ignore[union-attr]
-                required = True
-            elif annotation.canonical_path in {
-                "typing.NotRequired",
-                "typing_extensions.NotRequired",
-            }:
-                annotation = annotation.slice.elements[0]  # type: ignore[union-attr]
-                required = False
-        attrs.append(
-            _TypedDictAttr(
-                name=attr.name,
-                annotation=annotation,
-                docstring=attr.docstring,
-                required=required,
-            ),
-        )
+        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"] = attrs
-    return attrs
+    cls.extra["unpack_typeddict"]["_attributes"] = (required_attrs, optional_attrs)
+    return (required_attrs, optional_attrs)
 
 
-def _update_docstring(func: Function, attributes: Iterable[_TypedDictAttr], 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)
 
-    required = []
-    optional = []
-    for attribute in attributes:
-        if attribute["required"]:
-            required.append(attribute)
-        else:
-            optional.append(attribute)
-
     params_section = None
-    other_params_section = None
     sections = func.docstring.parsed
 
     # Find existing "Parameters" section.
     section_gen = (section for section in sections if section.kind is DocstringSectionKind.parameters)
     params_section = next(section_gen, None)
 
     # Pop original variadic keyword parameter from section.
-    varkw_param = None
     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:
-            varkw_param = params_section.value.pop(kwarg_pos)
+            params_section.value.pop(kwarg_pos)
 
     # If we have required parameters, add them to the "Parameters" section.
     if required:
@@ -118,55 +145,44 @@ def _update_docstring(func: Function, attributes: Iterable[_TypedDictAttr], kwpa
                 ),
             )
 
-        # Add back the original variadic keyword parameter if it was present,
-        # and if some parameters are optional.
-        if optional and varkw_param is not None:
-            params_section.value.append(
-                DocstringParameter(
-                    name=varkw_param.name,
-                    description=varkw_param.description,
-                    annotation=varkw_param.annotation,
-                ),
-            )
-
-    # If we have optional parameters, add them to the "Other parameters" section.
+    # If we have optional parameters, add them to the "Parameters" section too,
+    # with a default value of `...`.
     if optional:
-        # Create an "Other parameters" section if none exists.
-        section_gen = (section for section in sections if section.kind is DocstringSectionKind.other_parameters)
-        if (other_params_section := next(section_gen, None)) is None:
-            other_params_section = DocstringSectionOtherParameters([])
-            func.docstring.parsed.append(other_params_section)
+        # Create a "Parameters" section if none exists.
+        if params_section is None:
+            params_section = DocstringSectionParameters([])
+            func.docstring.parsed.append(params_section)
 
         # Add optional parameters to the section.
         for attr in optional:
-            other_params_section.value.append(
+            params_section.value.append(
                 DocstringParameter(
                     name=attr["name"],
                     description=attr["docstring"].value if attr["docstring"] else "",
                     annotation=attr["annotation"],
+                    value="...",
                 ),
             )
 
+    # TODO: Add `**kwargs` parameter if extra items are allowed.
 
-def _params_from_attrs(attrs: Iterable[_TypedDictAttr]) -> Parameters:
-    parameters = Parameters(Parameter(name="self", kind=ParameterKind.positional_or_keyword))
-    found_optional = False
-    for attr in attrs:
-        if attr["required"]:
-            parameters.add(
-                Parameter(
-                    name=attr["name"],
-                    annotation=attr["annotation"],
-                    kind=ParameterKind.keyword_only,
-                    docstring=attr["docstring"],
-                ),
-            )
-        else:
-            found_optional = True
-    if found_optional:
-        name = "_kwargs" if "kwargs" in parameters else "kwargs"
-        parameters.add(Parameter(name=name, kind=ParameterKind.var_keyword))
-    return parameters
+
+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):
@@ -181,18 +197,22 @@ def on_class(self, *, cls: Class, **kwargs: Any) -> None:  # noqa: ARG002
         else:
             return
 
-        attributes = _get_or_set_attrs(cls)
+        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, attributes)
+            _update_docstring(init, required, optional)
 
         # Remove attributes from the class, as they are now in the `__init__` method.
-        for attr in attributes:
+        for attr in chain(required, optional):
             cls.del_member(attr["name"])
 
     def on_function(self, *, func: Function, **kwargs: Any) -> None:  # noqa: ARG002
@@ -216,28 +236,22 @@ def on_function(self, *, func: Function, **kwargs: Any) -> None:  # noqa: ARG002
         else:
             return
 
-        attributes = _get_or_set_attrs(typed_dict)
-
-        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(attributes)
+        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, attributes, 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.