From c4712d22208ff39762cef75265071e0f5c9302cc Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Wed, 23 Sep 2026 15:01:44 +0200 Subject: [PATCH 1/3] perf(schema): compute optionally_keyed_by fields once per Schema class Schema.__post_init__ walked every annotation of the class and called get_origin() and get_args() on each of them to find the fields using optionally_keyed_by, every time an instance was created. The answer only depends on the class, so it is now computed once per class and cached. Creating an instance of a 10-field schema with two keyed-by fields takes 4.9us instead of 8.3us. In Firefox's task graph generation (38,223 tasks, `./mach taskgraph full -j 8`), the full task set is generated ~1-2s sooner (roughly 8%). --- src/taskgraph/util/schema.py | 49 ++++++++++++++++++++++++------------ test/test_util_schema.py | 41 ++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 16 deletions(-) diff --git a/src/taskgraph/util/schema.py b/src/taskgraph/util/schema.py index a78e74b5a..b75ee478d 100644 --- a/src/taskgraph/util/schema.py +++ b/src/taskgraph/util/schema.py @@ -6,6 +6,7 @@ import pprint import re import threading +import weakref from collections.abc import Mapping from typing import Annotated, Any, Literal, Optional, Union, get_args, get_origin @@ -336,6 +337,36 @@ def _caller_module_name(depth=1): return frame.f_globals.get("__name__", "schema") +_keyed_by_fields_cache = weakref.WeakKeyDictionary() + + +def _keyed_by_fields(cls): + """Return the (field name, OptionallyKeyedBy) pairs of the fields of `cls` + that use `optionally_keyed_by`. + + This only depends on the class, so it is computed once per class rather + than every time an instance is validated. + """ + try: + return _keyed_by_fields_cache[cls] + except KeyError: + pass + + 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) + ): + fields.append((field_name, args[1])) + + result = tuple(fields) + _keyed_by_fields_cache[cls] = result + return result + + class Schema( msgspec.Struct, kw_only=True, @@ -376,22 +407,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 _keyed_by_fields(type(self)): + keyed_by.validate(getattr(self, field_name)) # Validate mutually exclusive field groups. for group in getattr(self, "exclusive", []): diff --git a/test/test_util_schema.py b/test/test_util_schema.py index e10700b2b..612e7c5f6 100644 --- a/test/test_util_schema.py +++ b/test/test_util_schema.py @@ -2,7 +2,9 @@ # 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 gc import unittest +import weakref from typing import Optional import msgspec @@ -386,6 +388,45 @@ 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_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}}}) + + +def test_keyed_by_fields_cache_does_not_keep_classes_alive(): + S = Schema.from_dict({"field": optionally_keyed_by("foo", str, use_msgspec=True)}) + S.validate({"field": "a"}) + ref = weakref.ref(S) + + del S + gc.collect() + + assert ref() is None + + @pytest.mark.parametrize( "fields_dict, data, attr, expected", [ From bcccb508fb03eef56d56f2682ca44d9c772dd695 Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Wed, 23 Sep 2026 16:47:21 +0200 Subject: [PATCH 2/3] refactor(schema): store optionally_keyed_by fields on the Schema class --- src/taskgraph/util/schema.py | 48 ++++++++++++------------------------ test/test_util_schema.py | 13 ---------- 2 files changed, 16 insertions(+), 45 deletions(-) diff --git a/src/taskgraph/util/schema.py b/src/taskgraph/util/schema.py index b75ee478d..1e3c2173c 100644 --- a/src/taskgraph/util/schema.py +++ b/src/taskgraph/util/schema.py @@ -6,7 +6,6 @@ import pprint import re import threading -import weakref from collections.abc import Mapping from typing import Annotated, Any, Literal, Optional, Union, get_args, get_origin @@ -337,36 +336,6 @@ def _caller_module_name(depth=1): return frame.f_globals.get("__name__", "schema") -_keyed_by_fields_cache = weakref.WeakKeyDictionary() - - -def _keyed_by_fields(cls): - """Return the (field name, OptionallyKeyedBy) pairs of the fields of `cls` - that use `optionally_keyed_by`. - - This only depends on the class, so it is computed once per class rather - than every time an instance is validated. - """ - try: - return _keyed_by_fields_cache[cls] - except KeyError: - pass - - 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) - ): - fields.append((field_name, args[1])) - - result = tuple(fields) - _keyed_by_fields_cache[cls] = result - return result - - class Schema( msgspec.Struct, kw_only=True, @@ -394,11 +363,26 @@ 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 + # Find the fields that use `optionally_keyed_by` once per class, rather + # than every time an instance is validated. + 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.append((field_name, args[1])) + cls._keyed_by_fields = tuple(keyed_by_fields) + def __post_init__(self): if taskgraph.fast: return @@ -407,7 +391,7 @@ 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, keyed_by in _keyed_by_fields(type(self)): + for field_name, keyed_by in self._keyed_by_fields: keyed_by.validate(getattr(self, field_name)) # Validate mutually exclusive field groups. diff --git a/test/test_util_schema.py b/test/test_util_schema.py index 612e7c5f6..510928423 100644 --- a/test/test_util_schema.py +++ b/test/test_util_schema.py @@ -2,9 +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 gc import unittest -import weakref from typing import Optional import msgspec @@ -416,17 +414,6 @@ def test_optionally_keyed_by_from_dict(): S.validate({"field": {"by-foo": {"a": 1}}}) -def test_keyed_by_fields_cache_does_not_keep_classes_alive(): - S = Schema.from_dict({"field": optionally_keyed_by("foo", str, use_msgspec=True)}) - S.validate({"field": "a"}) - ref = weakref.ref(S) - - del S - gc.collect() - - assert ref() is None - - @pytest.mark.parametrize( "fields_dict, data, attr, expected", [ From f11aab54fbf8274aa0dc6fe5cc4219ef1eb02bb0 Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 10:58:02 +0200 Subject: [PATCH 3/3] fix(schema): validate optionally_keyed_by fields inherited from parent schemas --- src/taskgraph/util/schema.py | 12 +++++++----- test/test_util_schema.py | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/taskgraph/util/schema.py b/src/taskgraph/util/schema.py index 1e3c2173c..094a4d592 100644 --- a/src/taskgraph/util/schema.py +++ b/src/taskgraph/util/schema.py @@ -370,9 +370,9 @@ def __init_subclass__(cls, exclusive=None, **kwargs): if exclusive is not None: cls.exclusive = exclusive - # Find the fields that use `optionally_keyed_by` once per class, rather - # than every time an instance is validated. - keyed_by_fields = [] + 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 ( @@ -380,8 +380,10 @@ def __init_subclass__(cls, exclusive=None, **kwargs): and len(args) >= 2 and isinstance(args[1], OptionallyKeyedBy) ): - keyed_by_fields.append((field_name, args[1])) - cls._keyed_by_fields = tuple(keyed_by_fields) + 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: diff --git a/test/test_util_schema.py b/test/test_util_schema.py index 510928423..fece1d350 100644 --- a/test/test_util_schema.py +++ b/test/test_util_schema.py @@ -405,6 +405,24 @@ class IntSchema(Schema): 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)})