Remove Enum prefixes (#187)

Co-authored-by: nat <[email protected]> Co-authored-by: Tim Schmidt <[email protected]> Co-authored-by: Arun Babu Neelicattu <[email protected]>
danielgtaylor · Oct 25, 2023 · bd7de20 · bd7de20
1 parent d9b7608
commit bd7de20
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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",
+    }