`TypeAdapter.json_schema()` unable to render schema for custom `Annotated` type having a pydantic type in `PlainValidator.json_schema_input_type`

emaheuxPEREN commented 2 weeks ago

Initial Checks

[X] I confirm that I'm using Pydantic V2

Description

I define a custom Pydantic type with typing.Annotated + pydantic.PlainValidator(func, json_schema_input_type=OtherType) syntax :

Validation of this custom type works as expected
but I am unable to render its JSON schema with TypeAdapter.json_schema, as soon as OtherType is a pydantic type (e.g. a pydantic dataclass in my below example).

Full stack trace from example

{'properties': {'x': {'title': 'X', 'type': 'integer'}}, 'required': ['x'], 'title': 'MyNestedData', 'type': 'object'}  # OK

Traceback (most recent call last):
  File "test_pydantic_bug.py", line 29, in <module>
    print(TypeAdapter(MyRootData).json_schema())  # KeyError: '__main____MyNestedData-Input__1'
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/pydantic/type_adapter.py", line 135, in wrapped
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/pydantic/type_adapter.py", line 542, in json_schema
    return schema_generator_instance.generate(self.core_schema, mode=mode)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/pydantic/json_schema.py", line 416, in generate
    json_ref_counts = self.get_json_ref_counts(json_schema)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/pydantic/json_schema.py", line 2181, in get_json_ref_counts
    _add_json_refs(json_schema)
  File ".venv/lib/python3.11/site-packages/pydantic/json_schema.py", line 2170, in _add_json_refs
    _add_json_refs(self.definitions[defs_ref])
                   ~~~~~~~~~~~~~~~~^^^^^^^^^^
KeyError: '__main____MyNestedData-Input__1'

Example Code

from typing import Annotated, Self

from pydantic import TypeAdapter, PlainValidator
from pydantic.dataclasses import dataclass

@dataclass
class MyNestedData:
    x: int

print(TypeAdapter(MyNestedData).json_schema())  # OK

class _MyRootData:
    @classmethod
    def from_unsafe(cls, xxx) -> Self: ...

MyRootData = Annotated[
    _MyRootData,
    PlainValidator(_MyRootData.from_unsafe, json_schema_input_type=MyNestedData),
]

print(TypeAdapter(MyRootData).json_schema())  # KeyError: '__main____MyNestedData-Input__1'

Python, Pydantic & OS Version

pydantic version: 2.9.2
        pydantic-core version: 2.23.4
          pydantic-core build: profile=release pgo=false
                 install path: .venv/lib/python3.11/site-packages/pydantic
               python version: 3.11.9 (main, May  2 2024, 10:11:35) [GCC 12.2.0]
                     platform: Linux-6.1.0-26-amd64-x86_64-with-glibc2.36
             related packages: typing_extensions-4.12.2

sydney-runkle commented 2 weeks ago

Hi @emaheuxPEREN,

Looks like this is fixed on main, and you can confirm with v2.10.0b1. Let me know if you're still having issues, and I'd be happy to follow up!

emaheuxPEREN commented 2 weeks ago

Thank you! Indeed, the example do work with v2.10.0b1, but it won't work as soon as we use a generic container on top of the custom Pydantic type (which was my use case actually...):

MyRootData = Annotated[
    _MyRootData,
    PlainValidator(_MyRootData.from_unsafe, json_schema_input_type=list[MyNestedData]),  # MODIFIED here
]

print(TypeAdapter(MyRootData).json_schema())  # KeyError: '__main____MyNestedData-Input__1' (as previously)

using:

             pydantic version: 2.10.0b1
        pydantic-core version: 2.26.0

sydney-runkle commented 2 weeks ago

Ah yeah, looks like we're not resolving refs recursively... cc @viicos.

pydantic / pydantic