From fc4efd66bfbcc911da7bebc9b4d5765af1686633 Mon Sep 17 00:00:00 2001 From: fliiiix Date: Sat, 4 Nov 2023 10:43:28 +0100 Subject: [PATCH] [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. --- src/idl_gen_python.cpp | 13 +++++++++++-- tests/MyGame/Example/Any.py | 6 +++--- tests/MyGame/Example/AnyAmbiguousAliases.py | 6 +++--- tests/MyGame/Example/AnyUniqueAliases.py | 6 +++--- tests/MyGame/Example/NestedUnion/Any.py | 4 ++-- tests/monster_test_generated.py | 18 +++++++++--------- 6 files changed, 31 insertions(+), 22 deletions(-) diff --git a/src/idl_gen_python.cpp b/src/idl_gen_python.cpp index 7c5191836ed..0041df4c19e 100644 --- a/src/idl_gen_python.cpp +++ b/src/idl_gen_python.cpp @@ -591,12 +591,21 @@ class PythonStubGenerator { 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 @@ 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 index 998368952c3..17ce5503280 100644 --- 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 index 2225a036ca3..3a685f9a19b 100644 --- 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 index a3ba37e7fb2..837505e4055 100644 --- 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 index e0b85942414..caf272720d7 100644 --- a/tests/MyGame/Example/NestedUnion/Any.py +++ b/tests/MyGame/Example/NestedUnion/Any.py @@ -11,10 +11,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/monster_test_generated.py b/tests/monster_test_generated.py index f2b62cb2f2a..71139c65c06 100644 --- 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