Support for Python and enums
In the spirit and steps of https://grokbase.com/t/thrift/user/13614a6xd1/introspection-of-thrift-enums-in-python
(cherry picked from commit 4a8beb65f82a9525c5835cf93b664b7ac04ff076)
diff --git a/compiler/cpp/src/thrift/generate/t_py_generator.cc b/compiler/cpp/src/thrift/generate/t_py_generator.cc
index fe40fc2..af5821b 100644
--- a/compiler/cpp/src/thrift/generate/t_py_generator.cc
+++ b/compiler/cpp/src/thrift/generate/t_py_generator.cc
@@ -72,7 +72,9 @@
import_dynbase_ = "";
package_prefix_ = "";
for( iter = parsed_options.begin(); iter != parsed_options.end(); ++iter) {
- if( iter->first.compare("new_style") == 0) {
+ if( iter->first.compare("enum") == 0) {
+ gen_enum_ = true;
+ } else if( iter->first.compare("new_style") == 0) {
pwarning(0, "new_style is enabled by default, so the option will be removed in the near future.\n");
} else if( iter->first.compare("old_style") == 0) {
gen_newstyle_ = false;
@@ -295,6 +297,7 @@
* True if we should generate new-style classes.
*/
bool gen_newstyle_;
+ bool gen_enum_;
/**
* True if we should generate dynamic style classes.
@@ -463,6 +466,8 @@
<< "from thrift.protocol.TProtocol import TProtocolException"
<< endl
<< "from thrift.TRecursive import fix_spec"
+ << endl
+ << (gen_enum_ ? "from enum import IntEnum" : "")
<< endl;
if (gen_utf8strings_) {
@@ -502,9 +507,22 @@
*/
void t_py_generator::generate_enum(t_enum* tenum) {
std::ostringstream to_string_mapping, from_string_mapping;
+ std::string base_class;
- f_types_ << endl << endl << "class " << tenum->get_name() << (gen_newstyle_ ? "(object)" : "")
- << (gen_dynamic_ ? "(" + gen_dynbaseclass_ + ")" : "") << ":" << endl;
+ if (gen_enum_) {
+ base_class = "IntEnum";
+ } else if (gen_newstyle_) {
+ base_class = "object";
+ } else if (gen_dynamic_) {
+ base_class = gen_dynbaseclass_;
+ }
+
+ f_types_ << endl
+ << endl
+ << "class " << tenum->get_name()
+ << (base_class.empty() ? "" : "(" + base_class + ")")
+ << ":"
+ << endl;
indent_up();
generate_python_docstring(f_types_, tenum);
@@ -528,7 +546,9 @@
indent_down();
f_types_ << endl;
- f_types_ << to_string_mapping.str() << endl << from_string_mapping.str();
+ if (!gen_enum_) {
+ f_types_ << to_string_mapping.str() << endl << from_string_mapping.str();
+ }
}
/**
@@ -581,7 +601,14 @@
throw "compiler error: no const of base type " + t_base_type::t_base_name(tbase);
}
} else if (type->is_enum()) {
- out << value->get_integer();
+ out << indent();
+ int int_val = value->get_integer();
+ if (gen_enum_) {
+ t_enum_value* enum_val = ((t_enum*)type)->get_constant_by_value(int_val);
+ out << type->get_name() << "." << enum_val->get_name();
+ } else {
+ out << int_val;
+ }
} else if (type->is_struct() || type->is_xception()) {
out << type_name(type) << "(**{" << endl;
indent_up();
@@ -2216,7 +2243,7 @@
generate_deserialize_struct(out, (t_struct*)type, name);
} else if (type->is_container()) {
generate_deserialize_container(out, type, name);
- } else if (type->is_base_type() || type->is_enum()) {
+ } else if (type->is_base_type()) {
indent(out) << name << " = iprot.";
if (type->is_base_type()) {
@@ -2254,11 +2281,15 @@
default:
throw "compiler error: no Python name for base type " + t_base_type::t_base_name(tbase);
}
- } else if (type->is_enum()) {
- out << "readI32()";
}
out << endl;
-
+ } else if (type->is_enum()) {
+ if (gen_enum_) {
+ indent(out) << name << " = " << type_name(type) << "(iprot.readI32()).name";
+ } else {
+ out << "readI32()";
+ }
+ out << endl;
} else {
printf("DO NOT KNOW HOW TO DESERIALIZE FIELD '%s' TYPE '%s'\n",
tfield->get_name().c_str(),
@@ -2443,7 +2474,11 @@
throw "compiler error: no Python name for base type " + t_base_type::t_base_name(tbase);
}
} else if (type->is_enum()) {
- out << "writeI32(" << name << ")";
+ if (gen_enum_){
+ out << "writeI32(" << type_name(type) << "[" << name << "].value)";
+ } else {
+ out << "writeI32(" << name << ")";
+ }
}
out << endl;
} else {
@@ -2801,4 +2836,5 @@
" Add an import line to generated code to find the dynbase class.\n"
" package_prefix='top.package.'\n"
" Package prefix for generated files.\n"
- " old_style: Deprecated. Generate old-style classes.\n")
+ " old_style: Deprecated. Generate old-style classes.\n"
+ " enum: Generates Python's IntEnum, connects thrift to python enums . Python 3.4 and higher.\n")