Add enum option to typed_schema (#6546)

* Add enum option to typed_schema

* Assert keys all match
This commit is contained in:
Jesse Hills 2024-04-17 11:04:10 +12:00 committed by GitHub
parent 7d99676fe8
commit 6a1ea06744
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1590,6 +1590,10 @@ def typed_schema(schemas, **kwargs):
"""Create a schema that has a key to distinguish between schemas""" """Create a schema that has a key to distinguish between schemas"""
key = kwargs.pop("key", CONF_TYPE) key = kwargs.pop("key", CONF_TYPE)
default_schema_option = kwargs.pop("default_type", None) default_schema_option = kwargs.pop("default_type", None)
enum_mapping = kwargs.pop("enum", None)
if enum_mapping is not None:
assert isinstance(enum_mapping, dict)
assert set(enum_mapping.keys()) == set(schemas.keys())
key_validator = one_of(*schemas, **kwargs) key_validator = one_of(*schemas, **kwargs)
def validator(value): def validator(value):
@ -1600,6 +1604,9 @@ def typed_schema(schemas, **kwargs):
if schema_option is None: if schema_option is None:
raise Invalid(f"{key} not specified!") raise Invalid(f"{key} not specified!")
key_v = key_validator(schema_option) key_v = key_validator(schema_option)
if enum_mapping is not None:
key_v = add_class_to_obj(key_v, core.EnumValue)
key_v.enum_value = enum_mapping[key_v]
value = Schema(schemas[key_v])(value) value = Schema(schemas[key_v])(value)
value[key] = key_v value[key] = key_v
return value return value