diff --git a/src/taskgraph/util/schema.py b/src/taskgraph/util/schema.py index a78e74b5a..094a4d592 100644 --- a/src/taskgraph/util/schema.py +++ b/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", []): diff --git a/test/test_util_schema.py b/test/test_util_schema.py index e10700b2b..fece1d350 100644 --- a/test/test_util_schema.py +++ b/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", [