feat: #2101 Generate AsyncAPI schema for msgspec Structs

faststream/specification/asyncapi/message.py

@@ -11,23 +11,69 @@
     model_schema,
 )
 
+try:
+    import msgspec
+
+    HAS_MSGSPEC = True
+except ImportError:  # pragma: no cover
+    HAS_MSGSPEC = False
+
 if TYPE_CHECKING:
     from fast_depends.core import CallModel
 
 
+def is_msgspec_struct(annotation: Any) -> bool:
+    """Whether `annotation` is a msgspec Struct type."""
+    return HAS_MSGSPEC and isclass(annotation) and issubclass(annotation, msgspec.Struct)
+
+
+def get_msgspec_schema(struct: Any) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Build a JSON Schema for a msgspec Struct, shaped like Pydantic's.
+
+    Returns the struct's own schema inline plus the definitions any nested
+    structs live in, using the same `#/$defs/{name}` references Pydantic emits,
+    so the generators can hoist them into `components/schemas` unchanged.
+    """
+    (schema,), definitions = msgspec.json.schema_components(
+        [struct],
+        ref_template=f"#/{DEF_KEY}/{{name}}",
+    )
+    # A Struct is always emitted as a reference into the definitions, but stay
+    # defensive: an inline schema is still usable as-is.
+    name = schema.get("$ref", "").rsplit("/", 1)[-1]
+    body = dict(definitions.pop(name)) if name in definitions else dict(schema)
+    return body, definitions
+
+
 def parse_handler_params(call: "CallModel", prefix: str = "") -> dict[str, Any]:
     """Parses the handler parameters."""
     model_container = getattr(call, "serializer", call)
     model = cast("type[BaseModel] | None", getattr(model_container, "model", None))
     assert model
 
+    # Pydantic cannot build a schema for a msgspec Struct and raises rather than
+    # degrading, so structs are kept out of the model entirely and their schema
+    # is spliced back in by get_model_schema().
+    structs = {
+        p.field_name: p.field_type
+        for p in call.flat_params
+        if is_msgspec_struct(p.field_type)
+    }
+
     body = get_model_schema(
         create_model(
             model.__name__,
-            **{p.field_name: (p.field_type, p.default_value) for p in call.flat_params},  # type: ignore[call-overload]
+            **{  # type: ignore[call-overload]
+                p.field_name: (
+                    Any if p.field_name in structs else p.field_type,
+                    p.default_value,
+                )
+                for p in call.flat_params
+            },
         ),
         prefix=prefix,
         exclude=tuple(call.custom_fields.keys()),
+        structs=structs,
     )
 
     if body is None:
@@ -64,6 +110,7 @@ def get_model_schema(
     call: None,
     prefix: str = "",
     exclude: Sequence[str] = (),
+    structs: dict[str, Any] | None = None,
 ) -> None: ...
 
 
@@ -72,28 +119,52 @@ def get_model_schema(
     call: type[BaseModel],
     prefix: str = "",
     exclude: Sequence[str] = (),
+    structs: dict[str, Any] | None = None,
 ) -> dict[str, Any]: ...
 
 
 def get_model_schema(
     call: type[BaseModel] | None,
     prefix: str = "",
     exclude: Sequence[str] = (),
+    structs: dict[str, Any] | None = None,
 ) -> dict[str, Any] | None:
-    """Get the schema of a model."""
+    """Get the schema of a model.
+
+    `structs` maps parameter names to msgspec Structs that stand in the model as
+    `Any`, because Pydantic cannot describe them; their real schema is spliced
+    back in here.
+    """
     if call is None:
         return None
 
+    structs = dict(structs or {})
+
     params = {k: v for k, v in get_model_fields(call).items() if k not in exclude}
     params_number = len(params)
 
+    # Return annotations reach us as a model built elsewhere, with the Struct
+    # still on the field, so pick those up too.
+    for field_name, field in params.items():
+        if field_name not in structs and is_msgspec_struct(
+            getattr(field, "annotation", None),
+        ):
+            structs[field_name] = field.annotation
+
     if params_number == 0:
         return None
 
     model = None
     use_original_model = False
     if params_number == 1:
         name, param = next(iter(params.items()))
+        if name in structs:
+            # A lone Struct is described by its own schema, exactly like a lone
+            # Pydantic model is.
+            struct_body, definitions = get_msgspec_schema(structs[name])
+            if definitions:
+                struct_body[DEF_KEY] = definitions
+            return struct_body
         if (
             param.annotation
             and isclass(param.annotation)
@@ -107,6 +178,12 @@ def get_model_schema(
 
     body: dict[str, Any] = model_schema(model)
     body["properties"] = body.get("properties", {})
+    for param_name, struct in structs.items():
+        if param_name in body["properties"]:
+            struct_body, definitions = get_msgspec_schema(struct)
+            body["properties"][param_name] = struct_body
+            if definitions:
+                body.setdefault(DEF_KEY, {}).update(definitions)
     for i in exclude:
         body["properties"].pop(i, None)
     if required := body.get("required"):

tests/asyncapi/test_msgspec.py

@@ -0,0 +1,124 @@
+import msgspec
+import pytest
+from fast_depends.msgspec import MsgSpecSerializer
+from pydantic import BaseModel
+
+from faststream.nats import NatsBroker
+from faststream.specification import AsyncAPI
+
+VERSIONS = ("2.6.0", "3.0.0")
+
+
+class Address(msgspec.Struct):
+    city: str
+
+
+class User(msgspec.Struct):
+    name: str
+    age: int
+    address: Address
+
+
+def schemas(broker: NatsBroker, version: str) -> dict:
+    return (
+        AsyncAPI(broker, schema_version=version)
+        .to_specification()
+        .to_jsonable()["components"]["schemas"]
+    )
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_struct_argument_is_described_by_its_own_schema(version: str) -> None:
+    broker = NatsBroker()
+
+    @broker.subscriber("test")
+    async def handler(user: User) -> None: ...
+
+    assert schemas(broker, version)["User"] == {
+        "title": "User",
+        "type": "object",
+        "properties": {
+            "name": {"type": "string"},
+            "age": {"type": "integer"},
+            "address": {"$ref": "#/components/schemas/Address"},
+        },
+        "required": ["name", "age", "address"],
+    }
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_nested_struct_is_hoisted_into_components(version: str) -> None:
+    broker = NatsBroker()
+
+    @broker.subscriber("test")
+    async def handler(user: User) -> None: ...
+
+    # Like a nested Pydantic model, the nested Struct becomes its own component
+    # rather than being inlined or left as a dangling `#/$defs` reference.
+    assert schemas(broker, version)["Address"] == {
+        "title": "Address",
+        "type": "object",
+        "properties": {"city": {"type": "string"}},
+        "required": ["city"],
+    }
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_struct_alongside_other_arguments(version: str) -> None:
+    broker = NatsBroker()
+
+    @broker.subscriber("test")
+    async def handler(user: User, count: int) -> None: ...
+
+    payload = schemas(broker, version)["Handler:Message:Payload"]
+
+    assert payload["properties"]["count"] == {"title": "Count", "type": "integer"}
+    assert payload["properties"]["user"]["title"] == "User"
+    assert payload["required"] == ["user", "count"]
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_struct_return_annotation(version: str) -> None:
+    broker = NatsBroker()
+
+    @broker.publisher("out")
+    @broker.subscriber("in")
+    async def handler(count: int) -> Address: ...
+
+    assert schemas(broker, version)["Address"]["properties"] == {
+        "city": {"type": "string"}
+    }
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_struct_works_with_the_msgspec_serializer(version: str) -> None:
+    broker = NatsBroker(serializer=MsgSpecSerializer())
+
+    @broker.subscriber("test")
+    async def handler(user: User) -> None: ...
+
+    assert schemas(broker, version)["User"]["required"] == ["name", "age", "address"]
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_pydantic_models_are_unaffected(version: str) -> None:
+    class PydanticUser(BaseModel):
+        name: str
+
+    broker = NatsBroker()
+
+    @broker.subscriber("struct")
+    async def struct_handler(user: User) -> None: ...
+
+    @broker.subscriber("pydantic")
+    async def pydantic_handler(user: PydanticUser) -> None: ...
+
+    result = schemas(broker, version)
+
+    assert result["PydanticUser"] == {
+        "title": "PydanticUser",
+        "type": "object",
+        "properties": {"name": {"title": "Name", "type": "string"}},
+        "required": ["name"],
+    }
+    assert "User" in result