diff --git a/docs/source/flatc.md b/docs/source/flatc.md index 330ebc3d9a6..a0db6b5068b 100644 --- a/docs/source/flatc.md +++ b/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 - `--file-names-only` : Prints out files which would be generated by this command, one per diff --git a/include/flatbuffers/idl.h b/include/flatbuffers/idl.h index 4139529a11a..a8724f181f8 100644 --- a/include/flatbuffers/idl.h +++ b/include/flatbuffers/idl.h @@ -717,6 +717,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: @@ -860,6 +861,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), diff --git a/src/flatc.cpp b/src/flatc.cpp index cab0ee395a3..4e27b37f8ba 100644 --- a/src/flatc.cpp +++ b/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()); @@ -799,6 +804,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 " diff --git a/src/idl_gen_python.cpp b/src/idl_gen_python.cpp index 314c3c2d855..0a60bfdbf62 100644 --- a/src/idl_gen_python.cpp +++ b/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_typing && 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,14 @@ class PythonGenerator : public BaseGenerator { auto& field = **it; if (field.deprecated) continue; + // include import for enum type if used in this struct, we want type + // information, and we want modern enums. + if (IsEnum(field.value.type) && parser_.opts.python_typing && + parser_.opts.python_enum) { + imports.insert(ImportMapEntry{GenPackageReference(field.value.type), + namer_.Type(*field.value.type.enum_def)}); + } + GenStructAccessor(struct_def, field, code_ptr, imports); } @@ -1739,6 +1765,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,11 +1896,15 @@ class PythonGenerator : public BaseGenerator { break; } default: - // Scalar or sting fields. - field_type = GetBasePythonTypeForScalarAndString(base_type); - if (field.IsScalarOptional()) { - import_typing_list.insert("Optional"); - field_type = "Optional[" + field_type + "]"; + // Scalar or string fields. + if (parser_.opts.python_enum && IsEnum(field.value.type)) { + field_type = namer_.Type(*field.value.type.enum_def); + } else { + field_type = GetBasePythonTypeForScalarAndString(base_type); + if (field.IsScalarOptional()) { + import_typing_list.insert("Optional"); + field_type = "Optional[" + field_type + "]"; + } } break; } @@ -2647,6 +2682,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 +2801,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 +2817,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 +2829,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 +2894,50 @@ 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"; } } + + 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 =