perf(schema): compute optionally_keyed_by fields once per Schema class

src/taskgraph/util/schema.py

@@ -363,11 +363,28 @@ class MySchema(Schema, forbid_unknown_fields=False, kw_only=True):
             foo: str
     """
 
+    _keyed_by_fields = ()
+
     def __init_subclass__(cls, exclusive=None, **kwargs):
         super().__init_subclass__(**kwargs)
         if exclusive is not None:
             cls.exclusive = exclusive
 
+        keyed_by_fields = {}
+        for base in reversed(cls.__bases__):
+            keyed_by_fields.update(getattr(base, "_keyed_by_fields", ()))
+        for field_name, field_type in cls.__annotations__.items():
+            args = get_args(field_type)
+            if (
+                get_origin(field_type) is Annotated
+                and len(args) >= 2
+                and isinstance(args[1], OptionallyKeyedBy)
+            ):
+                keyed_by_fields[field_name] = args[1]
+            else:
+                keyed_by_fields.pop(field_name, None)
+        cls._keyed_by_fields = tuple(keyed_by_fields.items())
+
     def __post_init__(self):
         if taskgraph.fast:
             return
@@ -376,22 +393,8 @@ def __post_init__(self):
         # 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)
+        for field_name, keyed_by in self._keyed_by_fields:
+            keyed_by.validate(getattr(self, field_name))
 
         # Validate mutually exclusive field groups.
         for group in getattr(self, "exclusive", []):

test/test_util_schema.py

@@ -386,6 +386,52 @@ class TestSchema(Schema):
         TestSchema.validate({"field": {"by-foo": {"a": "b"}}})
 
 
+def test_optionally_keyed_by_per_class():
+    class StrSchema(Schema):
+        field: optionally_keyed_by("foo", str, use_msgspec=True)  # type: ignore
+
+    class IntSchema(Schema):
+        field: optionally_keyed_by("bar", int, use_msgspec=True)  # type: ignore
+        other: Optional[str] = None
+
+    for _ in range(2):
+        StrSchema.validate({"field": {"by-foo": {"a": "b"}}})
+        IntSchema.validate({"field": {"by-bar": {"a": 1}}})
+
+        with pytest.raises(msgspec.ValidationError):
+            StrSchema.validate({"field": {"by-foo": {"a": 1}}})
+
+        with pytest.raises(msgspec.ValidationError):
+            IntSchema.validate({"field": {"by-foo": {"a": 1}}})
+
+
+def test_optionally_keyed_by_subclass():
+    class BaseSchema(Schema, forbid_unknown_fields=False):
+        base: optionally_keyed_by("foo", str, use_msgspec=True)  # type: ignore
+
+    class SubSchema(BaseSchema):
+        sub: optionally_keyed_by("bar", int, use_msgspec=True)  # type: ignore
+
+    SubSchema.validate({"base": {"by-foo": {"a": "b"}}, "sub": {"by-bar": {"a": 1}}})
+
+    with pytest.raises(msgspec.ValidationError):
+        SubSchema.validate({"base": {"by-foo": {"a": 1}}, "sub": {"by-bar": {"a": 1}}})
+
+    with pytest.raises(msgspec.ValidationError):
+        SubSchema.validate(
+            {"base": {"by-foo": {"a": "b"}}, "sub": {"by-bar": {"a": "b"}}}
+        )
+
+
+def test_optionally_keyed_by_from_dict():
+    S = Schema.from_dict({"field": optionally_keyed_by("foo", str, use_msgspec=True)})
+
+    S.validate({"field": {"by-foo": {"a": "b"}}})
+
+    with pytest.raises(msgspec.ValidationError):
+        S.validate({"field": {"by-foo": {"a": 1}}})
+
+
 @pytest.mark.parametrize(
     "fields_dict, data, attr, expected",
     [