Repository: jcrist/msgspec
Pull request #345 (2023-03-19 20:27:50)

Support attrs types in json schema generation
This adds support for attrs types to `msgspec.json.schema` and `msgspec.inspect.type_info`. This was accidentally missed when support for encoding/decoding these types was added.


Patch:
diff --git a/msgspec/inspect.py b/msgspec/inspect.py
index f3b3683..71f0a54 100644
--- a/msgspec/inspect.py
+++ b/msgspec/inspect.py
@@ -23,6 +23,7 @@ from ._core import (
 from ._utils import (
     _CONCRETE_TYPES,
     _AnnotatedAlias,
+    get_dataclass_info as _get_dataclass_info,
     get_type_hints as _get_type_hints,
 )
 
@@ -496,7 +497,7 @@ class NamedTupleType(Type):
 
 
 class DataclassType(Type):
-    """A type corresponding to a `dataclasses.dataclass` type.
+    """A type corresponding to a `dataclasses` or `attrs` type.
 
     Parameters
     ----------
@@ -656,6 +657,10 @@ def _is_dataclass(t):
     return hasattr(t, "__dataclass_fields__")
 
 
+def _is_attrs(t):
+    return hasattr(t, "__attrs_attrs__")
+
+
 def _is_typeddict(t):
     try:
         return issubclass(t, dict) and hasattr(t, "__total__")
@@ -919,25 +924,22 @@ class _Translator:
                 for name, field_type in sorted(hints.items())
             )
             return out
-        elif _is_dataclass(t):
-            from dataclasses import MISSING, fields
-
+        elif _is_dataclass(t) or _is_attrs(t):
             if t in self.cache:
                 return self.cache[t]
             self.cache[t] = out = DataclassType(t, ())
-            hints = self._get_type_hints(t)
+            info, defaults, _, _ = _get_dataclass_info(t)
+            defaults = ((UNSET,) * (len(info) - len(defaults))) + defaults
             out.fields = tuple(
                 Field(
-                    name=f.name,
-                    encode_name=f.name,
-                    type=self.translate(hints[f.name]),
-                    required=f.default is MISSING and f.default_factory is MISSING,
-                    default=UNSET if f.default is MISSING else f.default,
-                    default_factory=(
-                        UNSET if f.default_factory is MISSING else f.default_factory
-                    ),
+                    name=name,
+                    encode_name=name,
+                    type=self.translate(typ),
+                    required=default is UNSET,
+                    default=UNSET if is_factory else default,
+                    default_factory=default if is_factory else UNSET,
                 )
-                for f in fields(t)
+                for (name, typ, is_factory), default in zip(info, defaults)
             )
             return out
         elif _is_namedtuple(t):


Test patch:
diff --git a/tests/test_inspect.py b/tests/test_inspect.py
index 41b2218..e135457 100644
--- a/tests/test_inspect.py
+++ b/tests/test_inspect.py
@@ -515,6 +515,26 @@ def test_dataclass():
     assert mi.type_info(Example) == sol
 
 
+def test_attrs():
+    attrs = pytest.importorskip("attrs")
+
+    @attrs.define
+    class Example:
+        x: int
+        y: int = 0
+        z: str = attrs.field(factory=str)
+
+    sol = mi.DataclassType(
+        Example,
+        fields=(
+            mi.Field("x", "x", mi.IntType()),
+            mi.Field("y", "y", mi.IntType(), required=False, default=0),
+            mi.Field("z", "z", mi.StrType(), required=False, default_factory=str),
+        ),
+    )
+    assert mi.type_info(Example) == sol
+
+
 @pytest.mark.parametrize("kind", ["struct", "namedtuple", "typeddict", "dataclass"])
 def test_self_referential_objects(kind):
     if kind == "struct":
diff --git a/tests/test_schema.py b/tests/test_schema.py
index df5b0d2..4409738 100644
--- a/tests/test_schema.py
+++ b/tests/test_schema.py
@@ -641,6 +641,58 @@ def test_dataclass():
     }
 
 
+def test_attrs():
+    attrs = pytest.importorskip("attrs")
+
+    @attrs.define
+    class Point:
+        x: int
+        y: int
+
+    @attrs.define
+    class Polygon:
+        """An example docstring"""
+
+        vertices: List[Point]
+        name: Union[str, None] = None
+        metadata: Dict[str, str] = attrs.field(factory=dict)
+
+    assert msgspec.json.schema(Polygon) == {
+        "$ref": "#/$defs/Polygon",
+        "$defs": {
+            "Polygon": {
+                "title": "Polygon",
+                "description": "An example docstring",
+                "type": "object",
+                "properties": {
+                    "vertices": {
+                        "type": "array",
+                        "items": {"$ref": "#/$defs/Point"},
+                    },
+                    "name": {
+                        "anyOf": [{"type": "string"}, {"type": "null"}],
+                        "default": None,
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "additionalProperties": {"type": "string"},
+                    },
+                },
+                "required": ["vertices"],
+            },
+            "Point": {
+                "title": "Point",
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                },
+                "required": ["x", "y"],
+            },
+        },
+    }
+
+
 @pytest.mark.parametrize("use_union_operator", [False, True])
 def test_union(use_union_operator):
     class Example(msgspec.Struct):