feat: support defining schemas as dicts

This provides an alternative way to define schemas, which IMO looks a
bit nicer than the class based approach. One big benefit of this method
is that you can use dashes in the keys, so it's possible to use both
underscores and dashes as identifiers and it's clear to the user which
is which.

Author: ahal

src/taskgraph/util/schema.py

@@ -2,6 +2,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import inspect
 import pprint
 import re
 import threading
@@ -318,6 +319,11 @@ def __getitem__(self, item):
         return self.schema[item]  # type: ignore
 
 
+def _caller_module_name(depth=1):
+    frame = inspect.stack()[depth + 1].frame
+    return frame.f_globals.get("__name__", "schema")
+
+
 class Schema(
     msgspec.Struct,
     kw_only=True,
@@ -345,6 +351,11 @@ class MySchema(Schema, forbid_unknown_fields=False, kw_only=True):
             foo: str
     """
 
+    def __init_subclass__(cls, exclusive=None, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if exclusive is not None:
+            cls.exclusive = exclusive
+
     def __post_init__(self):
         if taskgraph.fast:
             return
@@ -370,6 +381,76 @@ def __post_init__(self):
 
             keyed_by.validate(obj)
 
+        # Validate mutually exclusive field groups.
+        for group in getattr(self, "exclusive", []):
+            set_fields = [f for f in group if getattr(self, f) is not None]
+            if len(set_fields) > 1:
+                raise ValueError(
+                    f"{' and '.join(repr(f) for f in set_fields)} are mutually exclusive"
+                )
+
+    @classmethod
+    def from_dict(
+        cls,
+        fields_dict: dict[str, Any],
+        name: Optional[str] = None,
+        optional: bool = False,
+        **kwargs,
+    ) -> Union[type[msgspec.Struct], type[Optional[msgspec.Struct]]]:
+        """Create a Schema subclass dynamically from a dict of field definitions.
+
+        Each key is a field name and each value is either a type annotation or a
+        ``(type, default)`` tuple.  Fields typed as ``Optional[...]`` automatically
+        receive a default of ``None`` when no explicit default is provided.
+
+        Usage::
+
+            Schema.from_dict("MySchema", {
+                "required_field": str,
+                "optional_field": Optional[int],         # default None inferred
+                "explicit_default": (list[str], []),     # explicit default
+            })
+
+        Keyword arguments are forwarded to ``msgspec.defstruct`` (e.g.
+        ``forbid_unknown_fields=False``).
+        """
+        # Don't use `rename=kebab` by default as we can define kebab case
+        # properly in dicts.
+        kwargs.setdefault("rename", None)
+
+        # Ensure name and module are set correctly for error messages.
+        caller_module = _caller_module_name()
+        kwargs.setdefault("module", caller_module)
+        name = name or caller_module.rsplit(".", 1)[-1]
+
+        fields = []
+        for field_name, field_spec in fields_dict.items():
+            python_name = field_name.replace("-", "_")
+
+            if isinstance(field_spec, tuple):
+                typ, default = field_spec
+            else:
+                typ = field_spec
+                if get_origin(typ) is Union and type(None) in get_args(typ):
+                    default = None
+                else:
+                    default = msgspec.NODEFAULT
+
+            if field_name != python_name:
+                # Use msgspec.field to preserve the kebab-case encoded name.
+                # Explicit field names take priority over the struct-level rename.
+                fields.append(
+                    (python_name, typ, msgspec.field(name=field_name, default=default))
+                )
+            else:
+                fields.append((python_name, typ, default))
+
+        exclusive = kwargs.pop("exclusive", None)
+        result = msgspec.defstruct(name or "Schema", fields, bases=(cls,), **kwargs)
+        if exclusive:
+            result.exclusive = exclusive  # type: ignore[attr-defined]
+        return Optional[result] if optional else result  # type: ignore[valid-type]
+
     @classmethod
     def validate(cls, data):
         """Validate data against this schema."""

test/test_util_schema.py

@@ -3,6 +3,7 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import unittest
+from typing import Optional
 
 import msgspec
 import pytest
@@ -348,3 +349,53 @@ class TestSchema(Schema):
 
     with pytest.raises(msgspec.ValidationError):
         TestSchema.validate({"field": {"by-foo": {"a": "b"}}})
+
+
+@pytest.mark.parametrize(
+    "fields_dict, data, attr, expected",
+    [
+        ({"name": str}, {"name": "foo"}, "name", "foo"),
+        ({"count": Optional[int]}, {}, "count", None),
+        ({"tags": (list, [])}, {}, "tags", []),
+        ({"my-field": str}, {"my-field": "bar"}, "my_field", "bar"),
+    ],
+)
+def test_from_dict_valid(fields_dict, data, attr, expected):
+    S = Schema.from_dict(fields_dict)
+    result = msgspec.convert(data, S)
+    assert getattr(result, attr) == expected
+
+
+@pytest.mark.parametrize(
+    "fields_dict, data",
+    [
+        ({"name": str}, {}),
+        ({"my-field": str}, {"my_field": "bar"}),
+    ],
+)
+def test_from_dict_invalid(fields_dict, data):
+    S = Schema.from_dict(fields_dict)
+    with pytest.raises(msgspec.ValidationError):
+        msgspec.convert(data, S)
+
+
+@pytest.mark.parametrize(
+    "data, raises",
+    [
+        ({"a": "x", "b": "y"}, True),
+        ({"a": "x"}, False),
+        ({}, False),
+    ],
+)
+def test_exclusive(data, raises):
+    S = Schema.from_dict(
+        {"a": Optional[str], "b": Optional[str]},
+        exclusive=[["a", "b"]],
+    )
+    if raises:
+        with pytest.raises(
+            (ValueError, msgspec.ValidationError), match="mutually exclusive"
+        ):
+            msgspec.convert(data, S)
+    else:
+        msgspec.convert(data, S)