feat: allow docstring passon from typeddicts

Author: Anders Brams

openapi_python/generator/render.py

@@ -48,19 +48,26 @@ def _field_annotation(field: FieldDef) -> str:
     """
     Jinja2 filter to format a FieldDef's annotation.
     """
-    annotation = repr(_render_annotation(field.annotation))
+    annotation = repr(_render_field_annotation(field))
     if not field.required:
         annotation = f"NotRequired[{annotation}]"
     return annotation
 
 
 def _class_field_annotation(field: FieldDef, total_optional: bool) -> str:
-    annotation = _render_annotation(field.annotation)
+    annotation = _render_field_annotation(field)
     if not field.required and not total_optional:
         annotation = f"NotRequired[{annotation}]"
     return annotation
 
 
+def _render_field_annotation(field: FieldDef) -> str:
+    annotation = _render_annotation(field.annotation)
+    if field.description is None:
+        return annotation
+    return f"Annotated[{annotation}, _openapi_python_field({field.description!r})]"
+
+
 def _string_literal(value: str) -> str:
     return repr(value)
 
@@ -238,6 +245,10 @@ def _format_type_definition(defn: TypeAliasDef | TypedDictDef) -> str:
             return _format_typeddict(defn)
 
 
+def _has_field_descriptions(defns: tuple[TypedDictDef, ...]) -> bool:
+    return any(field.description is not None for defn in defns for field in defn.fields)
+
+
 def _call_parameters(op: OperationDef, *, generate_requests: bool) -> dict[str, str]:
     if not generate_requests:
         return {
@@ -395,6 +406,7 @@ def _render_types(
     ]
     return _render_template(
         "types.py.j2",
+        has_field_descriptions=_has_field_descriptions(typed_dicts),
         type_blocks="\n".join(blocks).strip() + "\n",
     )
 

openapi_python/generator/templates/types.py.j2

@@ -1,6 +1,21 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Any, Literal, NotRequired, TypeAlias, TypedDict
+{% if has_field_descriptions -%}
+from importlib import import_module
+{% endif -%}
+from typing import Annotated, Any, Literal, NotRequired, TypeAlias, TypedDict
+
+{% if has_field_descriptions %}
+def _openapi_python_field(description: str) -> object:
+    try:
+        # Allow FastAPI applications and other OpenAPI-spec-auto-generating tools 
+        # to use the same field description syntax as the generated client code 
+        # without a hard dependency on Pydantic.
+        field = import_module("pydantic").Field
+    except ImportError:
+        return description
+    return field(description=description)
+{% endif %}
 
 {{ type_blocks }}

tests/contract/docstrings/app.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Annotated, TypedDict
+
 from fastapi import FastAPI
 from pydantic import BaseModel, Field
 
@@ -12,6 +14,115 @@ class MyDTO(BaseModel):
     a: int = Field(description="This is the docstring for a single field")
 
 
+class CiscoAccessSwitchVLAN(TypedDict):
+    vlan_id: int
+
+
+class CiscoAccessSwitchVLANInterface(TypedDict):
+    name: str
+
+
+class CiscoAccessSwitchInterface(TypedDict):
+    name: str
+
+
+class ISEProfile(TypedDict):
+    name: str
+
+
+class CiscoAccessSwitchTemplateParams(TypedDict):
+    hostname: Annotated[str, Field(description="The hostname of the switch")]
+    vlans: Annotated[
+        list[CiscoAccessSwitchVLAN],
+        Field(
+            description=(
+                "All the VLANs that should be configured on the switch. These are "
+                "all the VLANs associated with the location in Nautobot."
+            )
+        ),
+    ]
+    vlan_interfaces: Annotated[
+        list[CiscoAccessSwitchVLANInterface],
+        Field(
+            description="All the virtual interfaces on the switch with the tag 'VLAN'."
+        ),
+    ]
+    client_interfaces: Annotated[
+        list[CiscoAccessSwitchInterface],
+        Field(description="All the client interfaces on the switch."),
+    ]
+    downlink_interfaces: Annotated[
+        list[CiscoAccessSwitchInterface],
+        Field(description="All the downlink interfaces on the switch."),
+    ]
+    uplink_interfaces: Annotated[
+        list[CiscoAccessSwitchInterface],
+        Field(description="All the uplink interfaces on the switch."),
+    ]
+    device_mgmt_ip: Annotated[
+        str,
+        Field(
+            description="The IP address assigned to the switch for management purposes."
+        ),
+    ]
+    management_vlan_id: Annotated[
+        int,
+        Field(
+            description=(
+                "The VLAN ID of the management VLAN. This is used for the "
+                "management interface and default gateway."
+            )
+        ),
+    ]
+    management_interface: Annotated[
+        str,
+        Field(
+            description=(
+                "The interface on which the management IP is configured. This is "
+                "used to determine which interface should be used for out-of-band "
+                "management access to the switch."
+            )
+        ),
+    ]
+    default_gateway: Annotated[
+        str,
+        Field(description="The default gateway for the management VLAN."),
+    ]
+    snmp_contact: Annotated[
+        str,
+        Field(description="The SNMP contact information for the switch."),
+    ]
+    snmp_location: Annotated[
+        str,
+        Field(description="The SNMP location information for the switch."),
+    ]
+    ise_profile: Annotated[
+        ISEProfile,
+        Field(description="The ISE profile to use for this switch."),
+    ]
+
+
 @app.get("/dto", response_model=MyDTO)
 def get_dto() -> MyDTO:
     return MyDTO(a=1)
+
+
+@app.get("/switch-template", response_model=CiscoAccessSwitchTemplateParams)
+def get_switch_template() -> CiscoAccessSwitchTemplateParams:
+    switch_vlan: CiscoAccessSwitchVLAN = {"vlan_id": 10}
+    switch_interface: CiscoAccessSwitchInterface = {"name": "GigabitEthernet1/0/1"}
+    return {
+        "hostname": "switch-01",
+        "vlans": [switch_vlan],
+        "vlan_interfaces": [{"name": "Vlan10"}],
+        "client_interfaces": [switch_interface],
+        "downlink_interfaces": [switch_interface],
+        "uplink_interfaces": [switch_interface],
+        "device_mgmt_ip": "192.0.2.10",
+        "management_vlan_id": 10,
+        "management_interface": "Vlan10",
+        "default_gateway": "192.0.2.1",
+        "snmp_contact": "Network Operations",
+        "snmp_location": "DC1",
+        "ise_profile": {"name": "default"},
+    }

tests/contract/docstrings/generate.py

@@ -6,9 +6,69 @@
 from pathlib import Path
 
 from app import app
+from fastapi import FastAPI
 
 from openapi_python.generator import GenerationRequest, generate_client
 
+SWITCH_TEMPLATE_FIELD_DOCSTRINGS = {
+    "hostname": "The hostname of the switch",
+    "vlans": (
+        "All the VLANs that should be configured on the switch. These are all "
+        "the VLANs associated with the location in Nautobot."
+    ),
+    "vlan_interfaces": "All the virtual interfaces on the switch with the tag 'VLAN'.",
+    "client_interfaces": "All the client interfaces on the switch.",
+    "downlink_interfaces": "All the downlink interfaces on the switch.",
+    "uplink_interfaces": "All the uplink interfaces on the switch.",
+    "device_mgmt_ip": "The IP address assigned to the switch for management purposes.",
+    "management_vlan_id": (
+        "The VLAN ID of the management VLAN. This is used for the management "
+        "interface and default gateway."
+    ),
+    "management_interface": (
+        "The interface on which the management IP is configured. This is used "
+        "to determine which interface should be used for out-of-band management "
+        "access to the switch."
+    ),
+    "default_gateway": "The default gateway for the management VLAN.",
+    "snmp_contact": "The SNMP contact information for the switch.",
+    "snmp_location": "The SNMP location information for the switch.",
+    "ise_profile": "The ISE profile to use for this switch.",
+}
+
+
+def _class_def(module: ast.Module, name: str) -> ast.ClassDef:
+    return next(
+        node
+        for node in module.body
+        if isinstance(node, ast.ClassDef) and node.name == name
+    )
+
+
+def _field_docstrings(class_def: ast.ClassDef) -> dict[str, str]:
+    docs: dict[str, str] = {}
+    for field, docstring in zip(class_def.body, class_def.body[1:], strict=False):
+        if not isinstance(field, ast.AnnAssign) or not isinstance(
+            field.target, ast.Name
+        ):
+            continue
+        if not isinstance(docstring, ast.Expr) or not isinstance(
+            docstring.value, ast.Constant
+        ):
+            continue
+        if isinstance(docstring.value.value, str):
+            docs[field.target.id] = docstring.value.value
+    return docs
+
+
+def _schema_property_descriptions(schema: dict) -> dict[str, str]:
+    properties = schema["properties"]
+    return {
+        name: prop["description"]
+        for name, prop in properties.items()
+        if "description" in prop
+    }
+
 
 def main() -> None:
     output_dir = Path(__file__).parent / "generated"
@@ -22,16 +82,26 @@ def main() -> None:
 
     source = (output_dir / "my_client" / "types.py").read_text()
     module = ast.parse(source)
-    dto = next(
-        node
-        for node in module.body
-        if isinstance(node, ast.ClassDef) and node.name == "MyDTO"
-    )
+    dto = _class_def(module, "MyDTO")
+    switch_template = _class_def(module, "CiscoAccessSwitchTemplateParams")
     generated_types = importlib.import_module("generated.my_client.types")
 
     assert ast.get_docstring(dto) == "This is a descriptive docstring"
     assert generated_types.MyDTO.__doc__ == "This is a descriptive docstring"
     assert "This is the docstring for a single field" in source
+    assert _field_docstrings(switch_template) == SWITCH_TEMPLATE_FIELD_DOCSTRINGS
+
+    api_b = FastAPI()
+
+    def get_switch_template() -> object:
+        return {}
+
+    api_b.get(
+        "/switch-template",
+        response_model=generated_types.CiscoAccessSwitchTemplateParams,
+    )(get_switch_template)
+    schema = api_b.openapi()["components"]["schemas"]["CiscoAccessSwitchTemplateParams"]
+    assert _schema_property_descriptions(schema) == SWITCH_TEMPLATE_FIELD_DOCSTRINGS
 
 
 if __name__ == "__main__":

tests/contract/docstrings/usage_async.py

@@ -3,7 +3,7 @@
 from typing import assert_type
 
 from generated.my_client import AsyncClient
-from generated.my_client.types import MyDTO
+from generated.my_client.types import CiscoAccessSwitchTemplateParams, MyDTO
 
 client = AsyncClient(base_url="http://testserver")
 
@@ -12,3 +12,8 @@ async def main() -> None:
     result = await client.get("/dto")()
     assert_type(result, MyDTO)
     assert_type(result["a"], int)
+
+    switch_template = await client.get("/switch-template")()
+    assert_type(switch_template, CiscoAccessSwitchTemplateParams)
+    assert_type(switch_template["hostname"], str)
+    assert_type(switch_template["management_vlan_id"], int)

tests/contract/docstrings/usage_sync.py

@@ -3,10 +3,15 @@
 from typing import assert_type
 
 from generated.my_client import Client
-from generated.my_client.types import MyDTO
+from generated.my_client.types import CiscoAccessSwitchTemplateParams, MyDTO
 
 client = Client(base_url="http://testserver")
 
 result = client.get("/dto")()
 assert_type(result, MyDTO)
 assert_type(result["a"], int)
+
+switch_template = client.get("/switch-template")()
+assert_type(switch_template, CiscoAccessSwitchTemplateParams)
+assert_type(switch_template["hostname"], str)
+assert_type(switch_template["management_vlan_id"], int)