fix: support optionally_keyed_by with underlying dict

Author: ahal

src/taskgraph/util/schema.py

@@ -7,7 +7,7 @@
 import threading
 from collections.abc import Mapping
 from functools import reduce
-from typing import Any, Literal, Optional, Union
+from typing import Annotated, Any, Literal, Optional, Union, get_args, get_origin
 
 import msgspec
 import voluptuous
@@ -70,11 +70,40 @@ def validate_schema(schema, obj, msg_prefix):
             raise Exception(f"{msg_prefix}\n{str(exc)}\n{pprint.pformat(obj)}")
 
 
-def UnionTypes(*types):
-    """Use `functools.reduce` to simulate `Union[*allowed_types]` on older
-    Python versions.
-    """
-    return reduce(lambda a, b: Union[a, b], types)
+class OptionallyKeyedBy:
+    """Metadata class for optionally_keyed_by fields in msgspec schemas."""
+
+    def __init__(self, *fields, wrapped_type):
+        self.fields = fields
+        self.wrapped_type = wrapped_type
+
+    @classmethod
+    def uses_keyed_by(cls, obj) -> bool:
+        if not isinstance(obj, dict) or len(obj) != 1:
+            return False
+
+        key = list(obj)[0]
+        if not key.startswith("by-"):
+            return False
+
+        return True
+
+    def validate(self, obj) -> None:
+        if not self.uses_keyed_by(obj):
+            # Not using keyed by, validate directly against wrapped type
+            msgspec.convert(obj, self.wrapped_type)
+            return
+
+        # First validate the outer keyed-by dict
+        bykeys = UnionTypes(*[Literal[f"by-{field}"] for field in self.fields])
+        msgspec.convert(obj, dict[bykeys, dict])
+
+        # Next validate each inner value. We call self.validate recursively to
+        # support nested `by-*` keys.
+        keyed_by_dict = list(obj.values())[0]
+        for value in keyed_by_dict.values():
+            self.validate(value)
+
 
 
 def optionally_keyed_by(*arguments, use_msgspec=False):
@@ -86,13 +115,15 @@ def optionally_keyed_by(*arguments, use_msgspec=False):
         use_msgspec: If True, return msgspec type hints; if False, return voluptuous validator
     """
     if use_msgspec:
-        # msgspec implementation - return type hints
+        # msgspec implementation - use Annotated[Any, OptionallyKeyedBy]
         _type = arguments[-1]
         if _type is object:
             return object
         fields = arguments[:-1]
-        bykeys = [Literal[f"by-{field}"] for field in fields]
-        return Union[_type, dict[UnionTypes(*bykeys), dict[str, Any]]]
+        wrapper = OptionallyKeyedBy(*fields, wrapped_type=_type)
+        # Annotating Any allows msgspec to accept any value without validation.
+        # The actual validation then happens in Schema.__post_init__
+        return Annotated[Any, wrapper]
     else:
         # voluptuous implementation - return validator function
         schema = arguments[-1]
@@ -291,6 +322,13 @@ def __getitem__(self, item):
         return self.schema[item]  # type: ignore
 
 
+def UnionTypes(*types):
+    """Use `functools.reduce` to simulate `Union[*allowed_types]` on older
+    Python versions.
+    """
+    return reduce(lambda a, b: Union[a, b], types)
+
+
 class Schema(
     msgspec.Struct,
     kw_only=True,
@@ -318,6 +356,31 @@ class MySchema(Schema, forbid_unknown_fields=False, kw_only=True):
             foo: str
     """
 
+    def __post_init__(self):
+        if taskgraph.fast:
+            return
+
+        # Validate fields that use optionally_keyed_by. We need to validate this
+        # manually because msgspec doesn't support union types with multiple
+        # dicts. Any fields that use `optionally_keyed_by("foo", dict)` would
+        # otherwise raise an exception.
+        for field_name, field_type in self.__class__.__annotations__.items():
+            origin = get_origin(field_type)
+            args = get_args(field_type)
+
+            if (
+                origin is not Annotated
+                or len(args) < 2
+                or not isinstance(args[1], OptionallyKeyedBy)
+            ):
+                # Not using `optionally_keyed_by`
+                continue
+
+            keyed_by = args[1]
+            obj = getattr(self, field_name)
+
+            keyed_by.validate(obj)
+
     @classmethod
     def validate(cls, data):
         """Validate data against this schema."""

test/test_util_schema.py

@@ -290,24 +290,36 @@ class TestSchema(Schema):
     TestSchema.validate({"field": "baz"})
     TestSchema.validate({"field": {"by-foo": {"a": "b", "c": "d"}}})
 
-    # Inner dict values are Any, so mixed types are accepted
-    TestSchema.validate({"field": {"by-foo": {"a": 1, "c": "d"}}})
+    with pytest.raises(msgspec.ValidationError):
+        TestSchema.validate({"field": 1})
+
+    with pytest.raises(msgspec.ValidationError):
+        TestSchema.validate({"field": {"by-bar": "a"}})
+
+    with pytest.raises(msgspec.ValidationError):
+        TestSchema.validate({"field": {"by-bar": {1: "b"}}})
 
     with pytest.raises(msgspec.ValidationError):
         TestSchema.validate({"field": {"by-bar": {"a": "b"}}})
 
+    with pytest.raises(msgspec.ValidationError):
+        TestSchema.validate({"field": {"by-foo": {"a": 1, "c": "d"}}})
+
 
-def test_optionally_keyed_by_mulitple_keys():
+def test_optionally_keyed_by_multiple_keys():
     class TestSchema(Schema):
         field: optionally_keyed_by("foo", "bar", str, use_msgspec=True)  # type: ignore
 
     TestSchema.validate({"field": {"by-foo": {"a": "b"}}})
     TestSchema.validate({"field": {"by-bar": {"x": "y"}}})
+    TestSchema.validate({"field": {"by-foo": {"a": {"by-bar": {"x": "y"}}}}})
 
     # Test invalid keyed-by field
     with pytest.raises(msgspec.ValidationError):
         TestSchema.validate({"field": {"by-unknown": {"a": "b"}}})
 
+    with pytest.raises(msgspec.ValidationError):
+        TestSchema.validate({"field": {"by-foo": {"a": {"by-bar": {"x": 1}}}}})
 
 def test_optionally_keyed_by_object_passthrough():
     """When the type argument is `object`, optionally_keyed_by returns object directly."""
@@ -320,14 +332,16 @@ def test_optionally_keyed_by_object_passthrough():
     assert msgspec.convert({"arbitrary": "dict"}, typ) == {"arbitrary": "dict"}
 
 
-@pytest.mark.xfail
 def test_optionally_keyed_by_dict():
     class TestSchema(Schema):
         field: optionally_keyed_by("foo", dict[str, str], use_msgspec=True)  # type: ignore
 
     TestSchema.validate({"field": {"by-foo": {"a": {"x": "y"}}}})
     TestSchema.validate({"field": {"a": "b"}})
 
+    with pytest.raises(msgspec.ValidationError):
+        TestSchema.validate({"field": {"a": 1}})
+
     with pytest.raises(msgspec.ValidationError):
         TestSchema.validate({"field": {"by-foo": {"a": {"x": 1}}}})