fix: loosen constraints on optionally_keyed_by for msgspec

Author: abhishekmadan30

src/taskgraph/util/schema.py

@@ -6,7 +6,7 @@
 import re
 from collections.abc import Mapping
 from functools import reduce
-from typing import Literal, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 import msgspec
 import voluptuous
@@ -85,9 +85,11 @@ def optionally_keyed_by(*arguments, use_msgspec=False):
     if use_msgspec:
         # msgspec implementation - return type hints
         _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, _type]]]
+        return Union[_type, dict[UnionTypes(*bykeys), dict[str, Any]]]
     else:
         # voluptuous implementation - return validator function
         schema = arguments[-1]

test/test_util_schema.py

@@ -290,8 +290,10 @@ def test_optionally_keyed_by():
         "by-foo": {"a": "b", "c": "d"}
     }
 
-    with pytest.raises(msgspec.ValidationError):
-        msgspec.convert({"by-foo": {"a": 1, "c": "d"}}, typ)
+    # Inner dict values are Any, so mixed types are accepted
+    assert msgspec.convert({"by-foo": {"a": 1, "c": "d"}}, typ) == {
+        "by-foo": {"a": 1, "c": "d"}
+    }
 
     with pytest.raises(msgspec.ValidationError):
         msgspec.convert({"by-bar": {"a": "b"}}, typ)
@@ -305,11 +307,22 @@ def test_optionally_keyed_by_mulitple_keys():
     }
     assert msgspec.convert({"by-bar": {"x": "y"}}, typ) == {"by-bar": {"x": "y"}}
 
-    with pytest.raises(msgspec.ValidationError):
-        msgspec.convert({"by-foo": {"a": 123, "c": "d"}}, typ)
-
-    with pytest.raises(msgspec.ValidationError):
-        msgspec.convert({"by-bar": {"a": 1}}, typ)
+    # Inner dict values are Any, so mixed types are accepted
+    assert msgspec.convert({"by-foo": {"a": 123, "c": "d"}}, typ) == {
+        "by-foo": {"a": 123, "c": "d"}
+    }
+    assert msgspec.convert({"by-bar": {"a": 1}}, typ) == {"by-bar": {"a": 1}}
 
     with pytest.raises(msgspec.ValidationError):
         msgspec.convert({"by-unknown": {"a": "b"}}, typ)
+
+
+def test_optionally_keyed_by_object_passthrough():
+    """When the type argument is `object`, optionally_keyed_by returns object directly."""
+    typ = optionally_keyed_by("foo", object, use_msgspec=True)
+    assert typ is object
+    # object accepts anything via msgspec.convert
+    assert msgspec.convert("hello", typ) == "hello"
+    assert msgspec.convert(42, typ) == 42
+    assert msgspec.convert({"by-foo": {"a": "b"}}, typ) == {"by-foo": {"a": "b"}}
+    assert msgspec.convert({"arbitrary": "dict"}, typ) == {"arbitrary": "dict"}