use_enum_values=True has no effect if another model with the same enum has use_enum_values=False

[vaultah]

Initial Checks

• I confirm that I'm using Pydantic V2

Description

If two separate models with different use_enum_values settings both have fields with the same enum, pydantic v2 doesn't use that enum's values in either model.

from pydantic import BaseModel, __version__ as pydantic_version
from enum import Enum


class MyEnum(str, Enum):
    A = 'a'


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


class OtherModel(BaseModel):
    also_my_enum: MyEnum


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


data = {'my_model': {'my_enum': 'a'}, 'other_model': {'also_my_enum': 'a'}}
print(f'{pydantic_version=}')
# pydantic_version='2.6.1'
print(Model.parse_obj(data))
# my_model=MyModel(my_enum=<MyEnum.A: 'a'>) other_model=OtherModel(also_my_enum=<MyEnum.A: 'a'>)

This is a regression from pydantic v1.

# pydantic_version='1.10.14'
# my_model=MyModel(my_enum='a') other_model=OtherModel(also_my_enum=<MyEnum.A: 'a'>)

To make MyModel receive enum's value in pydantic v2, I'd have to enable use_enum_values for OtherModel as well.

class OtherModel(BaseModel, use_enum_values=True):
    also_my_enum: MyEnum
...

# pydantic_version='2.6.1'
# my_model=MyModel(my_enum='a') other_model=OtherModel(also_my_enum='a')

Example Code

No response

Python, Pydantic & OS Version

pydantic version: 2.6.1
pydantic-core version: 2.16.2
  pydantic-core build: profile=release pgo=false
         install path: /too/much/unnecessary/information
       python version: 3.9.18 (main, Nov  2 2023, 16:51:22)  [Clang 14.0.3 (clang-1403.0.22.14.1)]
             platform: macOS-14.3-arm64-arm-64bit
     related packages: typing_extensions-4.9.0
               commit: unknown

[sydney-runkle]

@vaultah,

Yep, looks like a bug. Thanks for reporting the regression. You're welcome to submit a PR with a fix, if you'd like. I'll add this to the 2.7 milestone!

[NeevCohen]

@sydney-runkle I can try and look into this if it's still available 😄

[sydney-runkle]

@NeevCohen,

For sure! Go for it!

[NeevCohen]

Ok, I think I found the source of the problem.

When trying to debug this I wrote a test

def test_use_enum_values_in_single_submodel():
    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):
        other_model: OtherModel
        my_model: MyModel

    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)

When running this test, assert isinstance(obj.other_model.my_enum, MyEnum) would fail. If I changed the order of definition of the other_model and my_model fields in Model like so

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

Then assert not isinstance(obj.my_model.my_enum, MyEnum) would fail, so it seemed that the order of definition of theses fields was

After debugging a little, I saw that when generating schemas for MyModel and OtherModel, 2 correct schemas were created for MyEnum using the get_enum_core_schema, one which uses the enum values and another one which doesn't, like I would expect. But, both of the schemas for MyEnum had the same ref.

The problem seemed to be when generating the schema of Model, simplify_schema_references was called.

def simplify_schema_references(schema: core_schema.CoreSchema) -> core_schema.CoreSchema:  # noqa: C901
    definitions: dict[str, core_schema.CoreSchema] = {}
    ref_counts: dict[str, int] = defaultdict(int)
    involved_in_recursion: dict[str, bool] = {}
    current_recursion_ref_count: dict[str, int] = defaultdict(int)

    def collect_refs(s: core_schema.CoreSchema, recurse: Recurse) -> core_schema.CoreSchema:
        if s['type'] == 'definitions':
            for definition in s['definitions']:
                ref = get_ref(definition)
                assert ref is not None
                if ref not in definitions:
                    definitions[ref] = definition
                recurse(definition, collect_refs)
            return recurse(s['schema'], collect_refs)
        else:
            ref = get_ref(s)
            if ref is not None:
                new = recurse(s, collect_refs)
                new_ref = get_ref(new)
                if new_ref:
                    definitions[new_ref] = new
                return core_schema.definition_reference_schema(schema_ref=ref)
            else:
                return recurse(s, collect_refs)

As far as I understood, when simplify_schema_references was called with the schema of Model, while traversing the schema tree it found 2 schemas for MyEnum, one which uses the enum values and the other which wasn't, but since both of them had the same ref, the last one that was defined would be the one that took effect because of this line

if new_ref:
    definitions[new_ref] = new

which would basically overwrite the previous schema of MyEnum.

I hope that explanation makes sense.

Unfortunately, with my current limited understating of the core logic of the library, I couldn't think of a good solution for this problem. @sydney-runkle do you perhaps have an idea that could point me in the right direction?

[sydney-runkle]

@NeevCohen,

Great job looking into this. Funnily enough, I just spent some time battling with how simplify_schema_references handles various versions of schemas for the same object type (in my case, I'm currently trying to dig myself out of a hole with discriminated union logic).

I'll take a closer look at this today (maybe finding a solution to one of these problems will help us with the other). I'm guessing that we might need to do some better equality checks when we overwrite the ref values in the definitions dict.

Thanks again - I'll keep you updated!

[NeevCohen]

@sydney-runkle Awesome! Let me know if there's anything I can do to help out

[sydney-runkle]

Ah, @dmontagu came up with a fix: #8920