From 8eb450877e2f17272d3e0e94787e280190abea29 Mon Sep 17 00:00:00 2001
From: fliiiix <hi@l33t.name>
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                    | 8 +++++++-
 tests/MyGame/Example/NestedUnion/Any.py   | 3 ++-
 tests/MyGame/Example/NestedUnion/Color.py | 3 ++-
 tests/MyGame/Example/TestEnum.py          | 3 ++-
 4 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/idl_gen_python.cpp b/src/idl_gen_python.cpp
index a6bc0f51e8a8..35047f0a6947 100644
--- a/src/idl_gen_python.cpp
+++ b/src/idl_gen_python.cpp
@@ -108,7 +108,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.
diff --git a/tests/MyGame/Example/NestedUnion/Any.py b/tests/MyGame/Example/NestedUnion/Any.py
index e0b859424146..66d122c2c212 100644
--- 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
diff --git a/tests/MyGame/Example/NestedUnion/Color.py b/tests/MyGame/Example/NestedUnion/Color.py
index 6ba9ca1c6335..a9afc31e315d 100644
--- 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
index 4f5f60800b97..f3603f77bcda 100644
--- 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