[Python] Render enums as Python IntEnum

This allows enums to be type check with mypy. They will still behave like ints -> > IntEnum is the same as Enum, > but its members are also integers and can be used anywhere > that an integer can be used. > If any integer operation is performed with an IntEnum member, > the resulting value loses its enumeration status. https://docs.python.org/3/library/enum.html#enum.IntEnum Only if the --python-typing flag is set.
google · Jun 1, 2024 · 5cb1e2f · 5cb1e2f
1 parent 6ede1cc
commit 5cb1e2f
Show file tree

Hide file tree

Showing 8 changed files with 34 additions and 25 deletions.
diff --git a/src/idl_gen_python.cpp b/src/idl_gen_python.cpp
@@ -676,7 +676,13 @@ class PythonGenerator : public BaseGenerator {
   // Begin enum code with a class declaration.
   void BeginEnum(const EnumDef &enum_def, std::string *code_ptr) const {
     auto &code = *code_ptr;
-    code += "class " + namer_.Type(enum_def) + "(object):\n";
+    if (parser_.opts.python_typing) {
+      code += "from enum import IntEnum\n";
+      code += "class " + namer_.Type(enum_def) + "(IntEnum):\n";
+    }
+    else {
+      code += "class " + namer_.Type(enum_def) + "(object):\n";
+    }
   }
 
   // Starts a new line and then indents.
@@ -2432,7 +2438,7 @@ class PythonGenerator : public BaseGenerator {
     auto field_type = namer_.ObjectType(*ev.union_type.struct_def);
 
     code +=
-        GenIndents(1) + "if unionType == " + union_type + "()." + variant + ":";
+        GenIndents(1) + "if unionType == " + union_type + "." + variant + ":";
     if (parser_.opts.include_dependence_headers) {
       auto package_reference = GenPackageReference(ev.union_type);
       code += GenIndents(2) + "import " + package_reference;

diff --git a/tests/MyGame/Example/Any.py b/tests/MyGame/Example/Any.py
@@ -12,13 +12,13 @@ def AnyCreator(unionType, table):
     from flatbuffers.table import Table
     if not isinstance(table, Table):
         return None
-    if unionType == Any().Monster:
+    if unionType == Any.Monster:
         import MyGame.Example.Monster
         return MyGame.Example.Monster.MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == Any().TestSimpleTableWithEnum:
+    if unionType == Any.TestSimpleTableWithEnum:
         import MyGame.Example.TestSimpleTableWithEnum
         return MyGame.Example.TestSimpleTableWithEnum.TestSimpleTableWithEnumT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == Any().MyGame_Example2_Monster:
+    if unionType == Any.MyGame_Example2_Monster:
         import MyGame.Example2.Monster
         return MyGame.Example2.Monster.MonsterT.InitFromBuf(table.Bytes, table.Pos)
     return None
diff --git a/tests/MyGame/Example/AnyAmbiguousAliases.py b/tests/MyGame/Example/AnyAmbiguousAliases.py
@@ -12,13 +12,13 @@ def AnyAmbiguousAliasesCreator(unionType, table):
     from flatbuffers.table import Table
     if not isinstance(table, Table):
         return None
-    if unionType == AnyAmbiguousAliases().M1:
+    if unionType == AnyAmbiguousAliases.M1:
         import MyGame.Example.Monster
         return MyGame.Example.Monster.MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyAmbiguousAliases().M2:
+    if unionType == AnyAmbiguousAliases.M2:
         import MyGame.Example.Monster
         return MyGame.Example.Monster.MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyAmbiguousAliases().M3:
+    if unionType == AnyAmbiguousAliases.M3:
         import MyGame.Example.Monster
         return MyGame.Example.Monster.MonsterT.InitFromBuf(table.Bytes, table.Pos)
     return None
diff --git a/tests/MyGame/Example/AnyUniqueAliases.py b/tests/MyGame/Example/AnyUniqueAliases.py
@@ -12,13 +12,13 @@ def AnyUniqueAliasesCreator(unionType, table):
     from flatbuffers.table import Table
     if not isinstance(table, Table):
         return None
-    if unionType == AnyUniqueAliases().M:
+    if unionType == AnyUniqueAliases.M:
         import MyGame.Example.Monster
         return MyGame.Example.Monster.MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyUniqueAliases().TS:
+    if unionType == AnyUniqueAliases.TS:
         import MyGame.Example.TestSimpleTableWithEnum
         return MyGame.Example.TestSimpleTableWithEnum.TestSimpleTableWithEnumT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyUniqueAliases().M2:
+    if unionType == AnyUniqueAliases.M2:
         import MyGame.Example2.Monster
         return MyGame.Example2.Monster.MonsterT.InitFromBuf(table.Bytes, table.Pos)
     return None
diff --git a/tests/MyGame/Example/NestedUnion/Any.py b/tests/MyGame/Example/NestedUnion/Any.py
@@ -2,7 +2,8 @@
 
 # namespace: NestedUnion
 
-class Any(object):
+from enum import IntEnum
+class Any(IntEnum):
     NONE = 0
     Vec3 = 1
     TestSimpleTableWithEnum = 2
@@ -11,10 +12,10 @@ def AnyCreator(unionType, table):
     from flatbuffers.table import Table
     if not isinstance(table, Table):
         return None
-    if unionType == Any().Vec3:
+    if unionType == Any.Vec3:
         import MyGame.Example.NestedUnion.Vec3
         return MyGame.Example.NestedUnion.Vec3.Vec3T.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == Any().TestSimpleTableWithEnum:
+    if unionType == Any.TestSimpleTableWithEnum:
         import MyGame.Example.NestedUnion.TestSimpleTableWithEnum
         return MyGame.Example.NestedUnion.TestSimpleTableWithEnum.TestSimpleTableWithEnumT.InitFromBuf(table.Bytes, table.Pos)
     return None
diff --git a/tests/MyGame/Example/NestedUnion/Color.py b/tests/MyGame/Example/NestedUnion/Color.py
@@ -3,7 +3,8 @@
 # namespace: NestedUnion
 
 # Composite components of Monster color.
-class Color(object):
+from enum import IntEnum
+class Color(IntEnum):
     Red = 1
     # \brief color Green
     # Green is bit_flag with value (1u << 1)

diff --git a/tests/MyGame/Example/TestEnum.py b/tests/MyGame/Example/TestEnum.py
@@ -2,7 +2,8 @@
 
 # namespace: Example
 
-class TestEnum(object):
+from enum import IntEnum
+class TestEnum(IntEnum):
     A = 0
     B = 1
     C = 2
diff --git a/tests/monster_test_generated.py b/tests/monster_test_generated.py
@@ -39,11 +39,11 @@ def AnyCreator(unionType, table):
     from flatbuffers.table import Table
     if not isinstance(table, Table):
         return None
-    if unionType == Any().Monster:
+    if unionType == Any.Monster:
         return MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == Any().TestSimpleTableWithEnum:
+    if unionType == Any.TestSimpleTableWithEnum:
         return TestSimpleTableWithEnumT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == Any().MyGame_Example2_Monster:
+    if unionType == Any.MyGame_Example2_Monster:
         return MonsterT.InitFromBuf(table.Bytes, table.Pos)
     return None
 
@@ -58,11 +58,11 @@ def AnyUniqueAliasesCreator(unionType, table):
     from flatbuffers.table import Table
     if not isinstance(table, Table):
         return None
-    if unionType == AnyUniqueAliases().M:
+    if unionType == AnyUniqueAliases.M:
         return MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyUniqueAliases().TS:
+    if unionType == AnyUniqueAliases.TS:
         return TestSimpleTableWithEnumT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyUniqueAliases().M2:
+    if unionType == AnyUniqueAliases.M2:
         return MonsterT.InitFromBuf(table.Bytes, table.Pos)
     return None
 
@@ -77,11 +77,11 @@ def AnyAmbiguousAliasesCreator(unionType, table):
     from flatbuffers.table import Table
     if not isinstance(table, Table):
         return None
-    if unionType == AnyAmbiguousAliases().M1:
+    if unionType == AnyAmbiguousAliases.M1:
         return MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyAmbiguousAliases().M2:
+    if unionType == AnyAmbiguousAliases.M2:
         return MonsterT.InitFromBuf(table.Bytes, table.Pos)
-    if unionType == AnyAmbiguousAliases().M3:
+    if unionType == AnyAmbiguousAliases.M3:
         return MonsterT.InitFromBuf(table.Bytes, table.Pos)
     return None