[BUG] Backlinks cause errors when using documents as part of Pydantic union types
ldorigo opened this issue · 4 comments
ldorigo commented
Describe the bug
If I have a document containing backlinks; using that document within a pydantic basemodel that uses union types causes an exception. See simple MRO below.
To Reproduce
from typing import Literal
from uuid import uuid4
from beanie import BackLink
from beanie import Document
from beanie import Link
from beanie import init_beanie
from pydantic import UUID4
from pydantic import BaseModel
from pydantic import Field
class Improvement(Document):
id: UUID4 = Field(default_factory=uuid4)
issue_id: int
conversation_messages: list[Link["ConversationMessage"]] = []
class ConversationMessage(Document):
id: UUID4 = Field(default_factory=uuid4)
message_type: Literal["message"] # Placeholder for actual MessageType
improvements: list[BackLink["Improvement"]] = Field(
default=[], json_schema_extra={"original_field": "conversation_messages"}
)
class Settings:
is_root = True
class WelcomeMessage(ConversationMessage):
message_type: Literal["welcome_message"] = "welcome_message"
contents: str
class EscalationProposalMessage(ConversationMessage):
message_type: Literal["escalation_proposal"] = "escalation_proposal"
class Answer(BaseModel):
message: EscalationProposalMessage | WelcomeMessage
await init_beanie(
client.get_database(),
document_models=[Improvement, ConversationMessage, WelcomeMessage],
)
t = WelcomeMessage(contents="fwrwerF")
esc = EscalationProposalMessage()
Answer(message=t)
Note that the same happens with a simplified example without inheritance:
l
from uuid import uuid4
from beanie import BackLink
from beanie import Document
from beanie import Link
from beanie import init_beanie
from pydantic import UUID4
from pydantic import BaseModel
from pydantic import Field
class Improvement(Document):
id: UUID4 = Field(default_factory=uuid4)
issue_id: int
conversation_messages: list[Link["WelcomeMessage"] | Link["EscalationProposalMessage"]] = []
class WelcomeMessage(Document):
message_type: Literal["welcome_message"] = "welcome_message"
contents: str
improvements: list[BackLink["Improvement"]] = Field(
default=[], json_schema_extra={"original_field": "conversation_messages"}
)
class EscalationProposalMessage(Document):
message_type: Literal["escalation_proposal"] = "escalation_proposal"
improvements: list[BackLink["Improvement"]] = Field(
default=[], json_schema_extra={"original_field": "conversation_messages"}
)
class Answer(BaseModel):
message: EscalationProposalMessage | WelcomeMessage
await init_beanie(
client.get_database(),
document_models=[Improvement, WelcomeMessage, EscalationProposalMessage],
)
t = WelcomeMessage(contents="fwrwerF")
esc = EscalationProposalMessage()
Answer(message=t)
The error:
{
"name": "TypeError",
"message": "'WelcomeMessage' object does not support item assignment",
"stack": "---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[27], line 46
44 t = WelcomeMessage(contents=\"fwrwerF\")
45 esc = EscalationProposalMessage()
---> 46 Answer(message=t)
[... skipping hidden 1 frame]
File ~/.cache/pypoetry/virtualenvs/learnwise-chat-RkYLlhmr-py3.12/lib/python3.12/site-packages/beanie/odm/documents.py:240, in Document.fill_back_refs(cls, values)
238 @model_validator(mode=\"before\")
239 def fill_back_refs(cls, values):
--> 240 return cls._fill_back_refs(values)
File ~/.cache/pypoetry/virtualenvs/learnwise-chat-RkYLlhmr-py3.12/lib/python3.12/site-packages/beanie/odm/documents.py:229, in Document._fill_back_refs(cls, values)
221 values[field_name] = BackLink[link_info.document_class](
222 link_info.document_class
223 )
224 if (
225 link_info.link_type
226 in [LinkTypes.BACK_LIST, LinkTypes.OPTIONAL_BACK_LIST]
227 and field_name not in values
228 ):
--> 229 values[field_name] = [
230 BackLink[link_info.document_class](
231 link_info.document_class
232 )
233 ]
234 return values
TypeError: 'WelcomeMessage' object does not support item assignment"
}
Note that the error always happens on the second item in the union type: if we define Answer
with message: WelcomeMessage | EscalationProposalMessage
, then assigning a WelcomeMessage to messages
works, but it breaks when assigning an EscalationProposalMessage.
github-actions commented
This issue is stale because it has been open 30 days with no activity.
ldorigo commented
Not stale still an issue