codegen: Handle non-identifier characters in enum values

edgedb · May 26, 2023 · e1ec16d · e1ec16d
1 parent 297de72
commit e1ec16d
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 6 deletions.
diff --git a/edgedb/codegen/generator.py b/edgedb/codegen/generator.py
@@ -483,8 +483,8 @@ def _generate_code(
             buf = io.StringIO()
             self._imports.add("enum")
             print(f"class {rv}(enum.Enum):", file=buf)
-            for member in type_.members:
-                print(f'{INDENT}{member.upper()} = "{member}"', file=buf)
+            for member, member_id in self._to_unique_idents(type_.members):
+                print(f'{INDENT}{member_id.upper()} = "{member}"', file=buf)
             self._defs[rv] = buf.getvalue().strip()
 
         elif isinstance(type_, describe.RangeType):
@@ -537,3 +537,37 @@ def _snake_to_camel(self, name: str) -> str:
             return "".join(map(str.title, parts))
         else:
             return name
+
+    def _to_unique_idents(
+        self, names: typing.Iterable[typing.Tuple[str, str]]
+    ) -> typing.Iterator[str]:
+        dedup = set()
+        for name in names:
+            if name.isidentifier():
+                name_id = name
+                sep = name.endswith("_")
+            else:
+                sep = True
+                result = []
+                for i, c in enumerate(name):
+                    if c.isdigit():
+                        if i == 0:
+                            result.append("e_")
+                        result.append(c)
+                        sep = False
+                    elif c.isidentifier():
+                        result.append(c)
+                        sep = c == "_"
+                    elif not sep:
+                        result.append("_")
+                        sep = True
+                name_id = "".join(result)
+            rv = name_id
+            if not sep:
+                name_id = name_id + "_"
+            i = 1
+            while rv in dedup:
+                rv = f"{name_id}{i}"
+                i += 1
+            dedup.add(rv)
+            yield name, rv
diff --git a/tests/codegen/test-project2/generated_async_edgeql.py.assert b/tests/codegen/test-project2/generated_async_edgeql.py.assert
@@ -51,6 +51,9 @@ class LinkPropResultFriendsItem:
 class MyEnum(enum.Enum):
     THIS = "This"
     THAT = "That"
+    E_1 = "1"
+    F_B = "f. b"
+    F_B_1 = "f-b"
 
 
 @dataclasses.dataclass
@@ -218,7 +221,7 @@ async def my_query(
     return await executor.query_single(
         """\
         create scalar type MyScalar extending int64;
-        create scalar type MyEnum extending enum<This, That>;
+        create scalar type MyEnum extending enum<'This', 'That', '1', 'f. b', 'f-b'>;
 
         select {
         	a := <uuid>$a,

diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query.edgeql b/tests/codegen/test-project2/parpkg/subpkg/my_query.edgeql
@@ -1,5 +1,5 @@
 create scalar type MyScalar extending int64;
-create scalar type MyEnum extending enum<This, That>;
+create scalar type MyEnum extending enum<'This', 'That', '1', 'f. b', 'f-b'>;
 
 select {
 	a := <uuid>$a,

diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert
@@ -26,6 +26,9 @@ class NoPydanticValidation:
 class MyEnum(enum.Enum):
     THIS = "This"
     THAT = "That"
+    E_1 = "1"
+    F_B = "f. b"
+    F_B_1 = "f-b"
 
 
 @dataclasses.dataclass
@@ -144,7 +147,7 @@ async def my_query(
     return await executor.query_single(
         """\
         create scalar type MyScalar extending int64;
-        create scalar type MyEnum extending enum<This, That>;
+        create scalar type MyEnum extending enum<'This', 'That', '1', 'f. b', 'f-b'>;
 
         select {
         	a := <uuid>$a,

diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert b/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert
@@ -17,6 +17,9 @@ MyScalar = int
 class MyEnum(enum.Enum):
     THIS = "This"
     THAT = "That"
+    E_1 = "1"
+    F_B = "f. b"
+    F_B_1 = "f-b"
 
 
 @dataclasses.dataclass
@@ -135,7 +138,7 @@ def my_query(
     return executor.query_single(
         """\
         create scalar type MyScalar extending int64;
-        create scalar type MyEnum extending enum<This, That>;
+        create scalar type MyEnum extending enum<'This', 'That', '1', 'f. b', 'f-b'>;
 
         select {
         	a := <uuid>$a,