[Python] Render enums as Python IntEnum (#8145)
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.
diff --git a/src/idl_gen_python.cpp b/src/idl_gen_python.cpp
index 7c51918..0041df4 100644
--- a/src/idl_gen_python.cpp
+++ b/src/idl_gen_python.cpp
@@ -591,12 +591,21 @@
void GenerateEnumStub(std::stringstream &stub, const EnumDef *enum_def,
Imports *imports) const {
stub << "class " << namer_.Type(*enum_def);
- if (version_.major != 3) stub << "(object)";
+
+ if (version_.major == 3){
+ imports->Import("enum", "IntEnum");
+ stub << "(IntEnum)";
+ }
+ else {
+ stub << "(object)";
+ }
+
stub << ":\n";
for (const EnumVal *val : enum_def->Vals()) {
stub << " " << namer_.Variant(*val) << ": "
<< ScalarType(enum_def->underlying_type.base_type) << "\n";
}
+
if (parser_.opts.generate_object_based_api & enum_def->is_union) {
imports->Import("flatbuffers", "table");
stub << "def " << namer_.Function(*enum_def)
@@ -2432,7 +2441,7 @@
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
index 9983689..17ce550 100644
--- a/tests/MyGame/Example/Any.py
+++ b/tests/MyGame/Example/Any.py
@@ -12,13 +12,13 @@
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
index 2225a03..3a685f9 100644
--- a/tests/MyGame/Example/AnyAmbiguousAliases.py
+++ b/tests/MyGame/Example/AnyAmbiguousAliases.py
@@ -12,13 +12,13 @@
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
index a3ba37e..837505e 100644
--- a/tests/MyGame/Example/AnyUniqueAliases.py
+++ b/tests/MyGame/Example/AnyUniqueAliases.py
@@ -12,13 +12,13 @@
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
index e0b8594..caf2727 100644
--- a/tests/MyGame/Example/NestedUnion/Any.py
+++ b/tests/MyGame/Example/NestedUnion/Any.py
@@ -11,10 +11,10 @@
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/monster_test_generated.py b/tests/monster_test_generated.py
index f2b62cb..71139c6 100644
--- a/tests/monster_test_generated.py
+++ b/tests/monster_test_generated.py
@@ -39,11 +39,11 @@
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 @@
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 @@
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
