diff --git a/protobuf_to_pydantic/gen_model.py b/protobuf_to_pydantic/gen_model.py index f1f6fffe..5a2f34b3 100644 --- a/protobuf_to_pydantic/gen_model.py +++ b/protobuf_to_pydantic/gen_model.py @@ -224,6 +224,9 @@ def _get_field_info_dict_by_full_name(self, full_name: str) -> Optional["FieldIn return None return None + def _is_type_a_mapping(self, message_type: Descriptor) -> bool: + return message_type.name.endswith("Entry") and message_type.fields_by_name.keys() == ["key", "value"] + def _one_of_handle(self, descriptor: Descriptor) -> Tuple[Dict[str, "UseOneOfTypedDict"], Dict[str, Any]]: message_option_dict: "MessageOptionTypedDict" = self._message_option_dict.get( descriptor.name, {"message": {}, "one_of": {}, "nested": {}, "metadata": {}} @@ -297,7 +300,7 @@ def get_nested_message_dict_by_message(self, descriptor: Descriptor) -> Dict[str # nested support nested_message_dict: Dict[str, Type[Union[BaseModel, IntEnum]]] = {} for message in descriptor.nested_types: - if message.name.endswith("Entry"): + if self._is_type_a_mapping(message): continue nested_type: Any = self._parse_msg_to_pydantic_model(descriptor=message, root_descriptor=descriptor) nested_message_dict[message.full_name] = nested_type @@ -323,7 +326,7 @@ def _protobuf_field_type_is_type_message_handler(self, field_dataclass: FieldDat field_dataclass.field_default_factory = self._message_default_factory_dict_by_type_name[ protobuf_field.message_type.name ] - elif protobuf_field.message_type.name.endswith("Entry"): + elif self._is_type_a_mapping(protobuf_field.message_type): # support google.protobuf.MapEntry # key, value = column.message_type.fields field_dataclass.field_type_name = "map" @@ -442,7 +445,7 @@ def _protobuf_field_type_is_type_enum_handler(self, field_dataclass: FieldDataCl def _protobuf_field_lable_is_label_repeated_handler(self, field_dataclass: FieldDataClass) -> None: # support google.protobuf.array protobuf_field = field_dataclass.protobuf_field - if not (protobuf_field.message_type and protobuf_field.message_type.name.endswith("Entry")): + if not (protobuf_field.message_type and self._is_type_a_mapping(protobuf_field.message_type)): # I didn't know that Protobuf's Design of Maps and Lists would be so weird field_dataclass.field_type = List[field_dataclass.field_type] # type: ignore field_dataclass.field_default_factory = list