[Python] generate enums as IntEnum or IntFlag type when requested

docs/source/flatc.md

@@ -259,6 +259,8 @@ list of `FILES...`.
 
 -   `--python-typing` : Generate Python type annotations
 
+-   `--python-enum` : Generated enumerations inherit from `IntEnum` or `IntFlag` instead of `object`
+
 -   `--python-decode-obj-api-strings` : Decode bytes automaticaly with utf-8
 
 Additional gRPC options:

include/flatbuffers/idl.h

@@ -712,6 +712,7 @@ struct IDLOptions {
   /********************************** Python **********************************/
   bool python_no_type_prefix_suffix;
   bool python_typing;
+  bool python_enum;
   bool python_decode_obj_api_strings = false;
 
   // The target Python version. Can be one of the following:
@@ -855,6 +856,7 @@ struct IDLOptions {
         keep_proto_id(false),
         python_no_type_prefix_suffix(false),
         python_typing(false),
+        python_enum(false),
         python_gen_numpy(true),
         ts_omit_entrypoint(false),
         proto_id_gap_action(ProtoIdGapAction::WARNING),

src/flatc.cpp

@@ -262,6 +262,9 @@ const static FlatCOption flatc_options[] = {
     {"", "python-no-type-prefix-suffix", "",
      "Skip emission of Python functions that are prefixed with typenames"},
     {"", "python-typing", "", "Generate Python type annotations"},
+    {"", "python-enum", "",
+     "Generate enum types as IntEnum and IntFlag (assumes python-version >= "
+     "3"},
     {"", "python-version", "", "Generate code for the given Python version."},
     {"", "python-decode-obj-api-strings", "",
      "Decode bytes to strings for the Python Object API"},
@@ -694,6 +697,8 @@ FlatCOptions FlatCompiler::ParseFromCommandLineArguments(int argc,
         opts.python_no_type_prefix_suffix = true;
       } else if (arg == "--python-typing") {
         opts.python_typing = true;
+      } else if (arg == "--python-enum") {
+        opts.python_enum = true;
       } else if (arg.rfind("--python-version=", 0) == 0) {
         opts.python_version =
             arg.substr(std::string("--python-version=").size());
@@ -800,6 +805,11 @@ void FlatCompiler::ValidateOptions(const FlatCOptions& options) {
     Error("no options: specify at least one generator.", true);
   }
 
+  if (opts.python_enum &&
+      (opts.python_version.empty() || opts.python_version[0] != '3')) {
+    Error("--python-enum requires --python-version >= 3");
+  }
+
   if (opts.cs_gen_json_serializer && !opts.generate_object_based_api) {
     Error(
         "--cs-gen-json-serializer requires --gen-object-api to be set as "

src/idl_gen_python.cpp

@@ -612,9 +612,14 @@ class PythonStubGenerator {
 
     imports->Import("typing", "cast");
 
-    if (version_.major == 3) {
-      imports->Import("enum", "IntEnum");
-      stub << "(IntEnum)";
+    if (parser_.opts.python_enum) {
+      if (enum_def->attributes.Lookup("big_flags")) {
+        imports->Import("enum", "IntFlag");
+        stub << "(IntFlag)";
+      } else {
+        imports->Import("enum", "IntEnum");
+        stub << "(IntEnum)";
+      }
     } else {
       stub << "(object)";
     }
@@ -721,7 +726,20 @@ 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";
+
+    code += "class " + namer_.Type(enum_def);
+
+    if (parser_.opts.python_enum) {
+      if (enum_def.attributes.Lookup("bit_flags")) {
+        code += "(IntFlag)";
+      } else {
+        code += "(IntEnum)";
+      }
+    } else {
+      code += "(object)";
+    }
+
+    code += ":\n";
   }
 
   // Starts a new line and then indents.
@@ -1685,6 +1703,12 @@ class PythonGenerator : public BaseGenerator {
       auto& field = **it;
       if (field.deprecated) continue;
 
+      // include import for enum type if used in this struct
+      if (IsEnum(field.value.type)) {
+        imports.insert(ImportMapEntry{GenPackageReference(field.value.type),
+                                      namer_.Type(*field.value.type.enum_def)});
+      }
+
       GenStructAccessor(struct_def, field, code_ptr, imports);
     }
 
@@ -1739,6 +1763,11 @@ class PythonGenerator : public BaseGenerator {
     } else if (IsFloat(base_type)) {
       return float_const_gen_.GenFloatConstant(field);
     } else if (IsInteger(base_type)) {
+      // wrap the default value in the enum constructor to aid type hinting
+      if (parser_.opts.python_enum && IsEnum(field.value.type)) {
+        auto enum_type = namer_.Type(*field.value.type.enum_def);
+        return enum_type + "(" + field.value.constant + ")";
+      }
       return field.value.constant;
     } else {
       // For string, struct, and table.
@@ -1865,7 +1894,7 @@ class PythonGenerator : public BaseGenerator {
           break;
         }
         default:
-          // Scalar or sting fields.
+          // Scalar or string fields.
           field_type = GetBasePythonTypeForScalarAndString(base_type);
           if (field.IsScalarOptional()) {
             import_typing_list.insert("Optional");
@@ -2647,6 +2676,11 @@ class PythonGenerator : public BaseGenerator {
 
   std::string GenFieldTy(const FieldDef& field) const {
     if (IsScalar(field.value.type.base_type) || IsArray(field.value.type)) {
+      if (parser_.opts.python_enum) {
+        if (IsEnum(field.value.type)) {
+          return namer_.Type(*field.value.type.enum_def);
+        }
+      }
       const std::string ty = GenTypeBasic(field.value.type);
       if (ty.find("int") != std::string::npos) {
         return "int";
@@ -2761,7 +2795,8 @@ class PythonGenerator : public BaseGenerator {
   bool generate() {
     std::string one_file_code;
     ImportMap one_file_imports;
-    if (!generateEnums(&one_file_code)) return false;
+
+    if (!generateEnums(&one_file_code, one_file_imports)) return false;
     if (!generateStructs(&one_file_code, one_file_imports)) return false;
 
     if (parser_.opts.one_file) {
@@ -2776,7 +2811,8 @@ class PythonGenerator : public BaseGenerator {
   }
 
  private:
-  bool generateEnums(std::string* one_file_code) const {
+  bool generateEnums(std::string* one_file_code,
+                     ImportMap& one_file_imports) const {
     for (auto it = parser_.enums_.vec.begin(); it != parser_.enums_.vec.end();
          ++it) {
       auto& enum_def = **it;
@@ -2787,9 +2823,26 @@ class PythonGenerator : public BaseGenerator {
       }
 
       if (parser_.opts.one_file && !enumcode.empty()) {
+        if (parser_.opts.python_enum) {
+          if (enum_def.attributes.Lookup("bit_flags")) {
+            one_file_imports.insert({"enum", "IntFlag"});
+          } else {
+            one_file_imports.insert({"enum", "IntEnum"});
+          }
+        }
+
         *one_file_code += enumcode + "\n\n";
       } else {
         ImportMap imports;
+
+        if (parser_.opts.python_enum) {
+          if (enum_def.attributes.Lookup("bit_flags")) {
+            imports.insert({"enum", "IntFlag"});
+          } else {
+            imports.insert({"enum", "IntEnum"});
+          }
+        }
+
         const std::string mod =
             namer_.File(enum_def, SkipFile::SuffixAndExtension);
 
@@ -2835,49 +2888,52 @@ class PythonGenerator : public BaseGenerator {
   }
 
   // Begin by declaring namespace and imports.
-  void BeginFile(const std::string& name_space_name, const bool needs_imports,
-                 std::string* code_ptr, const std::string& mod,
-                 const ImportMap& imports) const {
+  void BeginFile(const std::string& name_space_name,
+                 const bool needs_default_imports, std::string* code_ptr,
+                 const std::string& mod, const ImportMap& imports) const {
     auto& code = *code_ptr;
     code = code + "# " + FlatBuffersGeneratedWarning() + "\n\n";
     code += "# namespace: " + name_space_name + "\n\n";
 
-    if (needs_imports) {
-      const std::string local_import = "." + mod;
-
+    if (needs_default_imports) {
       code += "import flatbuffers\n";
       if (parser_.opts.python_gen_numpy) {
         code += "from flatbuffers.compat import import_numpy\n";
       }
       if (parser_.opts.python_typing) {
         code += "from typing import Any\n";
-
-        for (auto import_entry : imports) {
-          // If we have a file called, say, "MyType.py" and in it we have a
-          // class "MyType", we can generate imports -- usually when we
-          // have a type that contains arrays of itself -- of the type
-          // "from .MyType import MyType", which Python can't resolve. So
-          // if we are trying to import ourself, we skip.
-          if (import_entry.first != local_import) {
-            code += "from " + import_entry.first + " import " +
-                    import_entry.second + "\n";
-          }
-        }
       }
-      if (parser_.opts.python_gen_numpy) {
-        code += "np = import_numpy()\n\n";
+    }
+    for (auto import_entry : imports) {
+      const std::string local_import = "." + mod;
+
+      // If we have a file called, say, "MyType.py" and in it we have a
+      // class "MyType", we can generate imports -- usually when we
+      // have a type that contains arrays of itself -- of the type
+      // "from .MyType import MyType", which Python can't resolve. So
+      // if we are trying to import ourself, we skip.
+      if (import_entry.first != local_import) {
+        code += "from " + import_entry.first + " import " +
+                import_entry.second + "\n";
       }
     }
+
+    code += "\n";
+
+    if (needs_default_imports && parser_.opts.python_gen_numpy) {
+      code += "np = import_numpy()\n\n";
+    }
   }
 
   // Save out the generated code for a Python Table type.
   bool SaveType(const std::string& defname, const Namespace& ns,
                 const std::string& classcode, const ImportMap& imports,
-                const std::string& mod, bool needs_imports) const {
+                const std::string& mod, bool needs_default_imports) const {
     if (classcode.empty()) return true;
 
     std::string code = "";
-    BeginFile(LastNamespacePart(ns), needs_imports, &code, mod, imports);
+    BeginFile(LastNamespacePart(ns), needs_default_imports, &code, mod,
+              imports);
     code += classcode;
 
     const std::string directories =

tests/PythonTest.sh

@@ -23,11 +23,18 @@ runtime_library_dir=${test_dir}/../python
 
 # Emit Python code for the example schema in the test dir:
 ${test_dir}/../flatc -p -o ${gen_code_path} -I include_test monster_test.fbs --gen-object-api
+${test_dir}/../flatc -p -o ${gen_code_path} -I include_test monster_test.fbs --gen-object-api --python-enum --python-version 3
 ${test_dir}/../flatc -p -o ${gen_code_path} -I include_test monster_test.fbs --gen-object-api --gen-onefile
 ${test_dir}/../flatc -p -o ${gen_code_path} -I include_test monster_extra.fbs --gen-object-api --python-typing --gen-compare
+${test_dir}/../flatc -p -o ${gen_code_path} -I include_test monster_extra.fbs --gen-object-api --python-typing --gen-compare --python-enum --python-version 3
 ${test_dir}/../flatc -p -o ${gen_code_path} -I include_test arrays_test.fbs --gen-object-api --python-typing
+${test_dir}/../flatc -p -o ${gen_code_path} -I include_test arrays_test.fbs --gen-object-api --python-enum --python-version 3
+${test_dir}/../flatc -p -o ${gen_code_path} -I include_test arrays_test.fbs --gen-object-api --python-enum --python-typing --python-version 3
+${test_dir}/../flatc -p -o ${gen_code_path} -I include_test arrays_test.fbs --gen-object-api --python-enum --python-typing --gen-onefile --python-version 3
 ${test_dir}/../flatc -p -o ${gen_code_path} -I include_test nested_union_test.fbs --gen-object-api --python-typing --python-decode-obj-api-strings
+${test_dir}/../flatc -p -o ${gen_code_path} -I include_test nested_union_test.fbs --gen-object-api --python-typing --python-decode-obj-api-strings --python-enum --python-version 3
 ${test_dir}/../flatc -p -o ${gen_code_path} -I include_test service_test.fbs --grpc --grpc-python-typed-handlers --python-typing --no-python-gen-numpy --gen-onefile
+${test_dir}/../flatc -p -o ${gen_code_path} -I include_test service_test.fbs --grpc --grpc-python-typed-handlers --python-typing --python-enum --no-python-gen-numpy --gen-onefile --python-version 3
 
 # Syntax: run_tests <interpreter> <benchmark vtable dedupes>
 #                   <benchmark read count> <benchmark build count>