Skip to content

Commit 2a4e0c1

Browse files
committed
fix: correctly handle optional enum generation
1 parent 1da8332 commit 2a4e0c1

File tree

1 file changed

+7
-3
lines changed

1 file changed

+7
-3
lines changed

protobuf_to_pydantic/gen_code.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -381,10 +381,14 @@ def _model_field_handle(self, model: Type[BaseModel], indent: int = 0) -> str:
381381
raise RuntimeError("can not load value type")
382382
# Type Hint handler
383383
if value_outer_type.__module__ != "builtins":
384-
if inspect.isclass(value_type) and issubclass(value_type, IntEnum):
385-
# Parse protobuf enum
384+
# Get the actual type to check, handling Optional cases
385+
type_to_check = value_type
386+
if isinstance(value_outer_type, _GenericAlias) and value_outer_type._name == "Optional":
387+
type_to_check = get_args(value_outer_type)[0]
388+
389+
if inspect.isclass(type_to_check) and issubclass(type_to_check, IntEnum):
386390
self._import_set.add("from enum import IntEnum")
387-
enum_code: str = self._gen_enum_py_code(value_type, indent=indent - self.code_indent)
391+
enum_code: str = self._gen_enum_py_code(type_to_check, indent=indent - self.code_indent)
388392
if enum_code:
389393
self._content_deque.append(enum_code)
390394
else:

0 commit comments

Comments
 (0)