fix: support optionally_keyed_by with underlying dict

Author: ahal

src/taskgraph/util/schema.py

@@ -6,7 +6,7 @@
 import re
 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
@@ -67,11 +67,12 @@ 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
 
 
 def optionally_keyed_by(*arguments, use_msgspec=False):
@@ -83,13 +84,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]
@@ -280,6 +283,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,
@@ -307,6 +317,48 @@ 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)
+            if not isinstance(obj, dict) or len(obj) != 1:
+                # Not using keyed by, validate direclty against wrapped type
+                msgspec.convert(obj, keyed_by.wrapped_type)
+                continue
+
+            key, keyed_by_dict = next(iter(obj.items()))
+            if not key.startswith("by-"):
+                # Not using keyed by, validate direclty against wrapped type
+                msgspec.convert(obj, keyed_by.wrapped_type)
+                continue
+
+            # First validate the outer keyed-by dict
+            bykeys = UnionTypes(*[Literal[f"by-{field}"] for field in keyed_by.fields])
+            msgspec.convert(obj, dict[bykeys, dict[str, Any]])
+
+            # Next validate each inner value against the wrapped type
+            for pattern, value in keyed_by_dict.items():
+                msgspec.convert(pattern, str)
+                msgspec.convert(value, keyed_by.wrapped_type)
+
     @classmethod
     def validate(cls, data):
         """Validate data against this schema."""

test/test_util_schema.py

@@ -290,6 +290,9 @@ class TestSchema(Schema):
     TestSchema.validate({"field": "baz"})
     TestSchema.validate({"field": {"by-foo": {"a": "b", "c": "d"}}})
 
+    with pytest.raises(msgspec.ValidationError):
+        TestSchema.validate({"field": 1})
+
     with pytest.raises(msgspec.ValidationError):
         TestSchema.validate({"field": {"by-bar": "a"}})
 
@@ -299,6 +302,9 @@ class TestSchema(Schema):
     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():
     class TestSchema(Schema):
@@ -323,14 +329,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}}}})