diff --git a/src/betterproto/casing.py b/src/betterproto/casing.py index 047f197ea..f7d0832b8 100644 --- a/src/betterproto/casing.py +++ b/src/betterproto/casing.py @@ -136,4 +136,8 @@ def lowercase_first(value: str) -> str: def sanitize_name(value: str) -> str: # https://www.python.org/dev/peps/pep-0008/#descriptive-naming-styles - return f"{value}_" if keyword.iskeyword(value) else value + if keyword.iskeyword(value): + return f"{value}_" + if not value.isidentifier(): + return f"_{value}" + return value diff --git a/src/betterproto/compile/naming.py b/src/betterproto/compile/naming.py index 1c2dbabee..baa9fc387 100644 --- a/src/betterproto/compile/naming.py +++ b/src/betterproto/compile/naming.py @@ -11,3 +11,11 @@ def pythonize_field_name(name: str) -> str: def pythonize_method_name(name: str) -> str: return casing.safe_snake_case(name) + + +def pythonize_enum_member_name(name: str, enum_name: str) -> str: + enum_name = casing.snake_case(enum_name).upper() + find = name.find(enum_name) + if find != -1: + name = name[find + len(enum_name) :].strip("_") + return casing.sanitize_name(name) diff --git a/src/betterproto/enum.py b/src/betterproto/enum.py index 1e3754326..529832251 100644 --- a/src/betterproto/enum.py +++ b/src/betterproto/enum.py @@ -57,7 +57,7 @@ def __new__( members = { name: value for name, value in namespace.items() - if not _is_descriptor(value) and name[0] != "_" + if not _is_descriptor(value) and not name.startswith("__") } cls = type.__new__( @@ -70,9 +70,6 @@ def __new__( # members become proper class variables for name, value in members.items(): - if _is_descriptor(value) or name[0] == "_": - continue - member = value_map.get(value) if member is None: member = cls.__new__(cls, name=name, value=value) # type: ignore diff --git a/src/betterproto/plugin/models.py b/src/betterproto/plugin/models.py index 1e22d8b3a..c7733816f 100644 --- a/src/betterproto/plugin/models.py +++ b/src/betterproto/plugin/models.py @@ -72,13 +72,13 @@ ) from betterproto.lib.google.protobuf.compiler import CodeGeneratorRequest -from ..casing import sanitize_name from ..compile.importing import ( get_type_reference, parse_source_type_name, ) from ..compile.naming import ( pythonize_class_name, + pythonize_enum_member_name, pythonize_field_name, pythonize_method_name, ) @@ -673,7 +673,9 @@ def __post_init__(self) -> None: # Get entries/allowed values for this Enum self.entries = [ self.EnumEntry( - name=sanitize_name(entry_proto_value.name), + name=pythonize_enum_member_name( + entry_proto_value.name, self.proto_obj.name + ), value=entry_proto_value.number, comment=get_comment( proto_file=self.source_file, path=self.path + [2, entry_number] diff --git a/tests/inputs/enum/enum.proto b/tests/inputs/enum/enum.proto index 97e12b4be..5e2e80c1f 100644 --- a/tests/inputs/enum/enum.proto +++ b/tests/inputs/enum/enum.proto @@ -15,3 +15,11 @@ enum Choice { FOUR = 4; THREE = 3; } + +// A "C" like enum with the enum name prefixed onto members, these should be stripped +enum ArithmeticOperator { + ARITHMETIC_OPERATOR_NONE = 0; + ARITHMETIC_OPERATOR_PLUS = 1; + ARITHMETIC_OPERATOR_MINUS = 2; + ARITHMETIC_OPERATOR_0_PREFIXED = 3; +} diff --git a/tests/inputs/enum/test_enum.py b/tests/inputs/enum/test_enum.py index 2aeb1d4f9..21a5ac3b9 100644 --- a/tests/inputs/enum/test_enum.py +++ b/tests/inputs/enum/test_enum.py @@ -1,4 +1,5 @@ from tests.output_betterproto.enum import ( + ArithmeticOperator, Choice, Test, ) @@ -102,3 +103,12 @@ def test_enum_mapped_on_parse(): # bonus: defaults after empty init are also mapped assert Test().choice.name == Choice.ZERO.name + + +def test_renamed_enum_members(): + assert set(ArithmeticOperator.__members__) == { + "NONE", + "PLUS", + "MINUS", + "_0_PREFIXED", + }