diff --git a/betterproto2/src/betterproto2/__init__.py b/betterproto2/src/betterproto2/__init__.py index 3a977e1..4b31b02 100644 --- a/betterproto2/src/betterproto2/__init__.py +++ b/betterproto2/src/betterproto2/__init__.py @@ -1019,7 +1019,13 @@ def parse(cls, data: bytes) -> Self: :class:`Message` The initialized message. """ + # Do not leave the parameter named ``data`` in the local frame while + # constructing a Pydantic dataclass. Generated modules commonly use + # a module-level ``data`` alias for nested enums; Pydantic's lazy + # forward-reference resolver can otherwise see this bytes argument + # and resolve ``data.SomeEnum`` against ``bytes``. with BytesIO(data) as stream: + del data return cls().load(stream) # For compatibility with other libraries. diff --git a/betterproto2/tests/test_pydantic_forward_refs.py b/betterproto2/tests/test_pydantic_forward_refs.py new file mode 100644 index 0000000..6cdeaee --- /dev/null +++ b/betterproto2/tests/test_pydantic_forward_refs.py @@ -0,0 +1,32 @@ +from enum import IntEnum +from types import SimpleNamespace + +from pydantic.dataclasses import dataclass + +import betterproto2 + + +@dataclass(eq=False, repr=False) +class GeneratedMessage(betterproto2.Message): + # This mirrors python-betterproto2 output where a direct child package is + # imported as ``data`` after the class declarations. + namespace: "data.GeneratedNamespace" = betterproto2.field( + 1, + betterproto2.TYPE_ENUM, + default_factory=lambda: data.GeneratedNamespace(0), + ) + + +class GeneratedNamespace(IntEnum): + UNKNOWN = 0 + + +data = SimpleNamespace(GeneratedNamespace=GeneratedNamespace) + + +def test_pydantic_forward_reference_resolves_on_cold_parse(): + assert GeneratedMessage.parse(b"").namespace is GeneratedNamespace.UNKNOWN + + +def test_pydantic_forward_reference_resolves_for_keyword_data_argument(): + assert GeneratedMessage.parse(data=b"").namespace is GeneratedNamespace.UNKNOWN