[BUG] Can't `Index` a `Link`

dedeswim commented 9 months ago

Describe the bug I would like to have a Link inside of a document, but I also want it to be unique, for which (in my understanding), I need Indexed. However, if I try that, I get an exception as below.

To Reproduce

from beanie import Document, Indexed, Link

class DocumentA(Document):
    ...

class DocumentB(Document):
    a: Indexed(Link[DocumentA], unique=True)

Expected behavior

The code should run without issues, however, I get the following exception:

Show exception

``` Traceback (most recent call last): File "/Users/edoardo/Library/Application Support/JetBrains/PyCharm2023.2/scratches/scratch.py", line 8, in class DocumentB(Document): File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_model_construction.py", line 178, in __new__ complete_model_class( File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_model_construction.py", line 468, in complete_model_class schema = cls.__get_pydantic_core_schema__(cls, handler) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/main.py", line 557, in __get_pydantic_core_schema__ return __handler(__source) ^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_schema_generation_shared.py", line 82, in __call__ schema = self._handler(__source_type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 448, in generate_schema schema = self._generate_schema(obj) ^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 689, in _generate_schema return self._model_schema(obj) ^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 519, in _model_schema {k: self._generate_md_field_schema(k, v, decorators) for k, v in fields.items()}, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 519, in {k: self._generate_md_field_schema(k, v, decorators) for k, v in fields.items()}, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 853, in _generate_md_field_schema common_field = self._common_field_schema(name, field_info, decorators) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 905, in _common_field_schema schema = self._apply_annotations( ^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 1610, in _apply_annotations schema = get_inner_schema(source_type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_schema_generation_shared.py", line 82, in __call__ schema = self._handler(__source_type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 1572, in inner_handler from_property = self._generate_schema_from_property(obj, obj) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 610, in _generate_schema_from_property schema = self._unpack_refs_defs(schema) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 567, in _unpack_refs_defs schema = flatten_schema_defs(schema) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 517, in flatten_schema_defs return _simplify_schema_references(schema, inline=False) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 432, in _simplify_schema_references schema = walk_core_schema(schema, collect_refs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 407, in walk_core_schema return f(schema, _dispatch) ^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 430, in collect_refs return recurse(s, collect_refs) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 183, in walk return f(schema.copy(), self._walk) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 430, in collect_refs return recurse(s, collect_refs) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 186, in _walk schema = self._schema_type_to_method[schema['type']](schema, f) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 195, in _handle_other_schemas schema['schema'] = self.walk(sub_schema, f) # type: ignore ^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 183, in walk return f(schema.copy(), self._walk) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 430, in collect_refs return recurse(s, collect_refs) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 186, in _walk schema = self._schema_type_to_method[schema['type']](schema, f) ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^ KeyError: 'Link' / ```

All the exceptions seem to be happening in Pydantic, but it I guess this is happening because of how Indexed is created?

Additional context

I am on beanie==1.23.1 and pydantic==2.4.2, the latest version for both.

It looks like this should've been fixed by #754 but apparently it is not?

dedeswim commented 9 months ago

It seems to be related to (if not the same as) #701. Feel free to close if it is actually the same.

roman-right commented 8 months ago

Hi! Thank you for the issue, I'll check it during the next bug-fixing session next week.

dedeswim commented 8 months ago

Fixed by #762 if using Annotated. The following now runs:

from typing import Annotated

from beanie import Document, Indexed, Link

class DocumentA(Document):
    ...

class DocumentB(Document):
    a: Annotated[Link[DocumentA], Indexed(unique=True)]

BeanieODM / beanie

[BUG] Can't `Index` a `Link` #763