KeyError during model_rebuild/recursive schemas

[commonism]

Initial Checks

• I confirm that I'm using Pydantic V2 installed directly from the main branch, or equivalent

Description

Defining a recursive model via model_rebuild sometimes fails with:

s = {'schema_ref': 'aiopenapi3.me.L8235d4c12_d58c_4316_829f_4f72825f8720_:94308158795872', 'type': 'definition-ref'}
recurse = <bound method _WalkCoreSchema._walk of <pydantic._internal._core_utils._WalkCoreSchema object at 0x7f8d34350a60>>

    def count_refs(s: core_schema.CoreSchema, recurse: Recurse) -> core_schema.CoreSchema:
        if not is_definition_ref_schema(s):
            return recurse(s, count_refs)
        ref = s['schema_ref']
        ref_counts[ref] += 1
        if current_recursion_ref_count[ref] != 0:
            involved_in_recursion[ref] = True
            return s
    
        current_recursion_ref_count[ref] += 1
>       recurse(all_defs[ref], count_refs)
E       KeyError: 'aiopenapi3.me.L8235d4c12_d58c_4316_829f_4f72825f8720_:94308158795872'

So far I can not reproduce easily, but I saw recent changes to _simplify_schema_references so … maybe … this is enough already?

Example Code

No response

Python, Pydantic & OS Version

pydantic version: 2.0a4
        pydantic-core version: 0.30.0 release build profile
                 install path: ~/venv/openapi3/lib/python3.10/site-packages/pydantic
               python version: 3.10.6 (main, Mar 10 2023, 10:55:28) [GCC 11.3.0]
                     platform: Linux-5.19.0-40-generic-x86_64-with-glibc2.35
     optional deps. installed: ['email-validator', 'typing-extensions']

[dmontagu]

@adriangb for context

@commonism I think it would be helpful if there was any way you could get us a code sample that could reproduce this issue.

[commonism]

Let's hope this allows to reproduce:

#!/bin/bash

rm -rf /tmp/keyerror /tmp/aiopenapi3
python3 -m venv /tmp/keyerror

git clone -b pydanticv2 https://github.com/commonism/aiopenapi3.git
cd aiopenapi3
/tmp/keyerror/bin/pip install --pre pydantic==2.0a4
/tmp/keyerror/bin/pip install .[tests]
/tmp/keyerror/bin/pytest tests/schema_test.py::test_schema_recursion

It'll use https://github.com/commonism/aiopenapi3/blob/master/tests/fixtures/schema-recursion.yaml

For me - it works when using a different file …

/tmp/keyerror/bin/aiopenapi3 validate tests/fixtures/schema-recursion.yaml

cp tests/fixtures/schema-recursion.yaml x.yaml
/tmp/keyerror/bin/aiopenapi3 validate x.yaml

smallest dd to reproduce:

openapi: "3.0.0"
info:
  version: 1.0.0
  title: Any OpenAPI schema with recursion

components:
  schemas:
    Expressions:
      type: array
      items:
        $ref: '#/components/schemas/Expression'

    Expression:
      type: object
      properties:
        Or:
          allOf:
            - $ref: '#/components/schemas/Expressions'
            - description: Return results that match either <code>Dimension</code> object.
        And:
          allOf:
            - $ref: '#/components/schemas/Expressions'
            - description: Return results that match both <code>Dimension</code> objects.
        Not:
          allOf:
            - $ref: '#/components/schemas/Expression'
            - description: Return results that don't match a <code>Dimension</code> object.

I never had it fail when using only two of the properties, sometimes it works with all three of them.

[commonism]

Do not waste time on this yet, I think this may be unrelated to pydantic itself but misuse of it.

[commonism]

I managed to reduce this example to something which may be acceptable.
Expect it to fail with something like:

    recurse(all_defs[ref], count_refs)
KeyError: '__main__.test.<locals>.allOfExpression_:94763011774496'

import random
from typing import List, ForwardRef
from pydantic import BaseModel

def test():
    class Expressions_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        items: List["__types['Expression']"]

    class Expression_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        Or: ForwardRef("__types['allOfExpressions']")
        And: ForwardRef("__types['allOfExpressions']")
        Not: ForwardRef("__types['allOfExpression']")

    class allOfExpression_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        Or: ForwardRef("__types['Expressions']")
        And: ForwardRef("__types['Expressions']")
        Not: ForwardRef("__types['Expression']")

    class allOfExpressions_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        items: List["__types['Expression']"]

    types = {"__types":{"Expression": Expression_, "Expressions": Expressions_, "allOfExpression": allOfExpression_, "allOfExpressions": allOfExpressions_}}


    l = [allOfExpressions_, Expressions_]#, allOfExpression_, Expression_]
#    l = list(types["__types"].values())
#    random.shuffle(l)
#    print(l)


    for i in l:
        i.model_rebuild(_types_namespace=types)

#for i in range(100):
#    test()
test()

If it works for you, use the commented lines to have it shuffle.

My understanding of the walker is limited, but I think this is related to L453 which ignores {'schema_ref': '…', 'type': 'definition-ref'}

pydantic/pydantic/_internal/_core_utils.py

Lines 445 to 456 in 5e81582

	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
	all_defs[ref] = recurse(definition, collect_refs).copy()
	return recurse(s['schema'], collect_refs)
	ref = get_ref(s)
	if ref is not None:
	all_defs[ref] = s
	return recurse(s, collect_refs)

[chriszs]

This is a +1 on a nested model from me. Not sure what I'm doing wrong.

[samuelcolvin]

@chriszs could you share your code?

[chriszs]

Reduced it down to:

from pydantic import BaseModel, ConfigDict

CONFIG = ConfigDict(undefined_types_warning=False)


class Division(BaseModel):
    model_config = CONFIG

    division_level: "DivisionLevel"


class DivisionLevel(BaseModel):
    model_config = CONFIG

    divisions: list["Division"]
    geometries: list["Geometry"]


class Geometry(BaseModel):
    model_config = CONFIG

    division_level: "DivisionLevel"
    points: list["Point"]


class Point(BaseModel):
    model_config = CONFIG

    geometry: "Geometry"


Point.model_rebuild()
DivisionLevel.model_rebuild()
Division.model_rebuild()
Geometry.model_rebuild()

Was getting:

  File "/pydantic/_internal/_core_utils.py", line 483, in count_refs
    recurse(all_defs[ref], count_refs)
KeyError: '__main__.Point:140665127664128'

I can fix it by flipping Division and DivisionLevel, so it definitely seems to be related to rebuild order:

DivisionLevel.model_rebuild()
Division.model_rebuild()

That works, but I wouldn't have expected that to be the error.

[commonism]

I can confirm.
Can be reduced to:

from pydantic import BaseModel, ConfigDict

CONFIG = ConfigDict(undefined_types_warning=False)

class Division(BaseModel):
    model_config = CONFIG
    division_level: "DivisionLevel"


class DivisionLevel(BaseModel):
    model_config = CONFIG

    division_level: "DivisionLevel"
    points: list["Point"]


class Point(BaseModel):
    model_config = CONFIG


DivisionLevel.model_rebuild()
Division.model_rebuild()

Any chance to have this addressed for beta #5919 ?

[dmontagu]

I definitely want to get this resolved ASAP, but I wasn't able to with a quick investigation. Hoping to find time to look into this more this week. Either way we'll definitely release a new release as soon as it is fixed

[commonism]

it changed and started working

In rev 2c26f54 & 6168b12

from typing import ForwardRef
from pydantic import BaseModel, ConfigDict

CONFIG = ConfigDict(undefined_types_warning=False)

class Division(BaseModel):
    model_config = CONFIG
    division_level: "DivisionLevel"


class DivisionLevel(BaseModel):
    model_config = CONFIG

    division_level: "DivisionLevel"
    points: list[ForwardRef("Point")] # difference is here


class Point(BaseModel):
    model_config = CONFIG


DivisionLevel.model_rebuild()
Division.model_rebuild()

works when using ForwardRef explicit

my example u works

[commonism]

current MRE

from typing import List, ForwardRef
from pydantic import BaseModel

def test():
    class Expressions_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        items: List["__types['Expression']"]

    class Expression_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        Or: ForwardRef("__types['allOfExpressions']")
        Not: ForwardRef("__types['allOfExpression']")

    class allOfExpression_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        Not: ForwardRef("__types['Expression']")

    class allOfExpressions_(BaseModel):
        model_config = dict(undefined_types_warning=False)
        items: List["__types['Expression']"]

    types = {"__types":{"Expression": Expression_, "Expressions": Expressions_, "allOfExpression": allOfExpression_, "allOfExpressions": allOfExpressions_}}


    l = [allOfExpressions_, Expressions_]
    for i in l:
        i.model_rebuild(_types_namespace=types)

test()

[lig]

It looks like the issue doesn't reproduce on the latest Pydantic main branch.

Here is a PR with the test added #6159

The test passes in on all supported Python versions.

@commonism does it works for you as well now?

[dmontagu]

I can confirm that all the example snippets in this issue currently don't error on main. I'm a bit concerned that since none of them attempt validation, the issue has been pushed off, but at the same time we (and by "we" I mostly mean @adriangb) have made many improvements so 🤞 it might just be fixed.

[dmontagu]

If anyone has suggestions of tests to add beyond what @lig has included in #6159 please suggest there

[commonism]

works for me