Skip to content

Commit

Permalink
Fix issue with config behavior being baked into the ref schema for en…
Browse files Browse the repository at this point in the history
…ums (#8920)

Co-authored-by: sydney-runkle <sydneymarierunkle@gmail.com>
  • Loading branch information
dmontagu and sydney-runkle committed Mar 1, 2024
1 parent 575f473 commit 2ff76ab
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 11 deletions.
20 changes: 9 additions & 11 deletions pydantic/_internal/_std_types_schema.py
Expand Up @@ -81,28 +81,19 @@ def get_json_schema(_, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
# so that we can still generate a valid json schema.
return core_schema.is_instance_schema(enum_type, metadata={'pydantic_js_functions': [get_json_schema]})

use_enum_values = config.get('use_enum_values', False)

if len(cases) == 1:
expected = repr(cases[0].value)
else:
expected = ', '.join([repr(case.value) for case in cases[:-1]]) + f' or {cases[-1].value!r}'

def to_enum(__input_value: Any) -> Enum:
try:
enum_field = enum_type(__input_value)
if use_enum_values:
return enum_field.value
return enum_field
return enum_type(__input_value)
except ValueError:
# The type: ignore on the next line is to ignore the requirement of LiteralString
raise PydanticCustomError('enum', f'Input should be {expected}', {'expected': expected}) # type: ignore

strict_python_schema = core_schema.is_instance_schema(enum_type)
if use_enum_values:
strict_python_schema = core_schema.chain_schema(
[strict_python_schema, core_schema.no_info_plain_validator_function(lambda x: x.value)]
)

to_enum_validator = core_schema.no_info_plain_validator_function(to_enum)
if issubclass(enum_type, int):
Expand Down Expand Up @@ -132,9 +123,16 @@ def to_enum(__input_value: Any) -> Enum:
else:
lax = to_enum_validator
strict = core_schema.json_or_python_schema(json_schema=to_enum_validator, python_schema=strict_python_schema)
return core_schema.lax_or_strict_schema(
enum_schema = core_schema.lax_or_strict_schema(
lax_schema=lax, strict_schema=strict, ref=enum_ref, metadata={'pydantic_js_functions': [get_json_schema]}
)
use_enum_values = config.get('use_enum_values', False)
if use_enum_values:
enum_schema = core_schema.chain_schema(
[enum_schema, core_schema.no_info_plain_validator_function(lambda x: x.value)]
)

return enum_schema


@dataclasses.dataclass(**slots_true)
Expand Down
19 changes: 19 additions & 0 deletions tests/test_types.py
Expand Up @@ -6410,3 +6410,22 @@ def test_on_error_omit_top_level() -> None:
# if it hits the top level, but this documents the current behavior at least
with pytest.raises(SchemaError, match='Uncaught Omit error'):
ta.validate_python('a')


def test_diff_enums_diff_configs() -> None:
class MyEnum(str, Enum):
A = 'a'

class MyModel(BaseModel, use_enum_values=True):
my_enum: MyEnum

class OtherModel(BaseModel):
my_enum: MyEnum

class Model(BaseModel):
my_model: MyModel
other_model: OtherModel

obj = Model.model_validate({'my_model': {'my_enum': 'a'}, 'other_model': {'my_enum': 'a'}})
assert not isinstance(obj.my_model.my_enum, MyEnum)
assert isinstance(obj.other_model.my_enum, MyEnum)

0 comments on commit 2ff76ab

Please sign in to comment.