Skip to content

Commit

Permalink
Python3 type hints
Browse files Browse the repository at this point in the history 
args not optional type hint

UUID generator support
  • Loading branch information
@arkuhn
Konstantin Pozdniakov authored and arkuhn committed Feb 16, 2024
1 parent 50787cc commit 819977b
Showing 1 changed file with 126 additions and 7 deletions.
133 changes: 126 additions & 7 deletions compiler/cpp/src/thrift/generate/t_py_generator.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ class t_py_generator : public t_generator {
gen_twisted_ = false;
gen_dynamic_ = false;
gen_enum_ = false;
gen_type_hints_ = false;
coding_ = "";
gen_dynbaseclass_ = "";
gen_dynbaseclass_exc_ = "";
Expand Down Expand Up @@ -127,6 +128,9 @@ class t_py_generator : public t_generator {
gen_tornado_ = true;
} else if( iter->first.compare("coding") == 0) {
coding_ = iter->second;
} else if( iter->first.compare("type_hints") == 0) {
gen_type_hints_ = true;
gen_enum_ = true;
} else {
throw "unknown option py:" + iter->first;
}
Expand Down Expand Up @@ -254,12 +258,16 @@ class t_py_generator : public t_generator {
std::string declare_argument(t_field* tfield);
std::string render_field_default_value(t_field* tfield);
std::string type_name(t_type* ttype);
std::string function_signature(t_function* tfunction, bool interface = false);
std::string function_signature(t_function* tfunction, bool interface = false, bool send_part = false);
std::string argument_list(t_struct* tstruct,
std::vector<std::string>* pre = nullptr,
std::vector<std::string>* post = nullptr);
std::string type_to_enum(t_type* ttype);
std::string type_to_spec_args(t_type* ttype);
std::string type_to_py_type(t_type* type);
std::string member_hint(t_type* type, t_field::e_req req);
std::string arg_hint(t_type* type);
std::string func_hint(t_type* type);

static bool is_valid_namespace(const std::string& sub_namespace) {
return sub_namespace == "twisted";
Expand Down Expand Up @@ -316,6 +324,11 @@ class t_py_generator : public t_generator {

bool gen_slots_;

/**
* True if we should generate classes type hints and type checks in write methods.
*/
bool gen_type_hints_;

std::string copy_options_;

/**
Expand Down Expand Up @@ -462,12 +475,19 @@ string t_py_generator::py_autogen_comment() {
*/
string t_py_generator::py_imports() {
ostringstream ss;
if (gen_type_hints_) {
ss << "from __future__ import annotations" << endl
<< "import typing" << endl;
}

ss << "from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, "
"TApplicationException"
<< endl
<< "from thrift.protocol.TProtocol import TProtocolException"
<< endl
<< "from thrift.TRecursive import fix_spec"
<< endl
<< "from uuid import UUID"
<< endl;
if (gen_enum_) {
ss << "from enum import IntEnum" << endl;
Expand Down Expand Up @@ -525,6 +545,7 @@ void t_py_generator::generate_enum(t_enum* tenum) {
<< (base_class.empty() ? "" : "(" + base_class + ")")
<< ":"
<< endl;

indent_up();
generate_python_docstring(f_types_, tenum);

Expand Down Expand Up @@ -599,6 +620,9 @@ string t_py_generator::render_const_value(t_type* type, t_const_value* value) {
out << emit_double_as_string(value->get_double());
}
break;
case t_base_type::TYPE_UUID:
out << "UUID(\"" << get_escaped_string(value) << "\")";
break;
default:
throw "compiler error: no const of base type " + t_base_type::t_base_name(tbase);
}
Expand Down Expand Up @@ -884,7 +908,8 @@ void t_py_generator::generate_py_struct_definition(ostream& out,
<< "'] = " << (*m_iter)->get_name() << endl;
}
} else {
indent(out) << "self." << (*m_iter)->get_name() << " = " << (*m_iter)->get_name() << endl;
indent(out) << "self." << (*m_iter)->get_name() << member_hint((*m_iter)->get_type(), (*m_iter)->get_req())
<< " = " << (*m_iter)->get_name() << endl;
}
}

Expand Down Expand Up @@ -1155,7 +1180,7 @@ void t_py_generator::generate_py_struct_writer(ostream& out, t_struct* tstruct)

indent(out) << "def write(self, oprot):" << endl;
indent_up();

indent(out) << "self.validate()" << endl;
indent(out) << "if oprot._fast_encode is not None and self.thrift_spec is not None:" << endl;
indent_up();

Expand Down Expand Up @@ -1562,7 +1587,7 @@ void t_py_generator::generate_service_client(t_service* tservice) {
}

f_service_ << endl;
indent(f_service_) << "def send_" << function_signature(*f_iter, false) << ":" << endl;
indent(f_service_) << "def send_" << function_signature(*f_iter, false, true) << ":" << endl;
indent_up();

std::string argsname = (*f_iter)->get_name() + "_args";
Expand Down Expand Up @@ -2341,6 +2366,9 @@ void t_py_generator::generate_deserialize_field(ostream& out,
case t_base_type::TYPE_DOUBLE:
out << "readDouble()";
break;
case t_base_type::TYPE_UUID:
out << "readUuid()";
break;
default:
throw "compiler error: no Python name for base type " + t_base_type::t_base_name(tbase);
}
Expand Down Expand Up @@ -2533,6 +2561,9 @@ void t_py_generator::generate_serialize_field(ostream& out, t_field* tfield, str
case t_base_type::TYPE_DOUBLE:
out << "writeDouble(" << name << ")";
break;
case t_base_type::TYPE_UUID:
out << "writeUuid(" << name << ")";
break;
default:
throw "compiler error: no Python name for base type " + t_base_type::t_base_name(tbase);
}
Expand All @@ -2559,6 +2590,11 @@ void t_py_generator::generate_serialize_field(ostream& out, t_field* tfield, str
* @param prefix String prefix to attach to all fields
*/
void t_py_generator::generate_serialize_struct(ostream& out, t_struct* tstruct, string prefix) {
if(gen_type_hints_) {
indent(out) << "if not type(" << prefix << ") is " << type_to_py_type(tstruct)
<< ": raise ValueError('" << prefix << " is not " << type_to_py_type(tstruct) << "')" << endl;
}

(void)tstruct;
indent(out) << prefix << ".write(oprot)" << endl;
}
Expand Down Expand Up @@ -2706,12 +2742,16 @@ void t_py_generator::generate_python_docstring(ostream& out, t_doc* tdoc) {
*/
string t_py_generator::declare_argument(t_field* tfield) {
std::ostringstream result;
result << tfield->get_name() << "=";
t_field::e_req req = tfield->get_req();
result << tfield->get_name() << member_hint(tfield->get_type(), req);

result << " = ";
if (tfield->get_value() != nullptr) {
result << render_field_default_value(tfield);
} else {
result << "None";
}

return result.str();
}

Expand All @@ -2735,7 +2775,7 @@ string t_py_generator::render_field_default_value(t_field* tfield) {
* @param tfunction Function definition
* @return String of rendered function definition
*/
string t_py_generator::function_signature(t_function* tfunction, bool interface) {
string t_py_generator::function_signature(t_function* tfunction, bool interface, bool send_part) {
vector<string> pre;
vector<string> post;
string signature = tfunction->get_name() + "(";
Expand All @@ -2745,6 +2785,10 @@ string t_py_generator::function_signature(t_function* tfunction, bool interface)
}

signature += argument_list(tfunction->get_arglist(), &pre, &post) + ")";
if (!send_part) {
signature += func_hint(tfunction->get_returntype());
}

return signature;
}

Expand Down Expand Up @@ -2775,6 +2819,7 @@ string t_py_generator::argument_list(t_struct* tstruct, vector<string>* pre, vec
result += ", ";
}
result += (*f_iter)->get_name();
result += arg_hint((*f_iter)->get_type());
}
if (post) {
for (s_iter = post->begin(); s_iter != post->end(); ++s_iter) {
Expand Down Expand Up @@ -2804,8 +2849,79 @@ string t_py_generator::type_name(t_type* ttype) {
return ttype->get_name();
}

string t_py_generator::arg_hint(t_type* type) {
if (gen_type_hints_) {
return ": " + type_to_py_type(type);
}

return "";
}

string t_py_generator::member_hint(t_type* type, t_field::e_req req) {
if (gen_type_hints_) {
if (req != t_field::T_REQUIRED) {
return ": typing.Optional[" + type_to_py_type(type) + "]";
}
else {
return ": " + type_to_py_type(type);
}
}

return "";
}

string t_py_generator::func_hint(t_type* type) {
if (gen_type_hints_) {
return " -> " + type_to_py_type(type);
}

return "";
}

/**
* Converts the parse type to a Python type
*/
string t_py_generator::type_to_py_type(t_type* type) {
type = get_true_type(type);

if (type->is_binary()) {
return "bytes";
}
if (type->is_base_type()) {
t_base_type::t_base tbase = ((t_base_type*)type)->get_base();
switch (tbase) {
case t_base_type::TYPE_VOID:
return "None";
case t_base_type::TYPE_STRING:
return "str";
case t_base_type::TYPE_BOOL:
return "bool";
case t_base_type::TYPE_I8:
case t_base_type::TYPE_I16:
case t_base_type::TYPE_I32:
case t_base_type::TYPE_I64:
return "int";
case t_base_type::TYPE_DOUBLE:
return "float";
case t_base_type::TYPE_UUID:
return "UUID";
}
} else if (type->is_enum() || type->is_struct() || type->is_xception()) {
return type_name(type);
} else if (type->is_map()) {
return "dict[" + type_to_py_type(((t_map*)type)->get_key_type()) + ", " + type_to_py_type(((t_map*)type)->get_val_type()) + "]";
} else if (type->is_set()) {
return "set[" + type_to_py_type(((t_set*)type)->get_elem_type()) + "]";
} else if (type->is_list()) {
return "list[" + type_to_py_type(((t_list*)type)->get_elem_type()) + "]";
}

throw "INVALID TYPE IN type_to_py_type: " + type->get_name();
}


/**
* Converts the parse type to a Python tyoe
* Converts the parse type to a Python type enum
*/
string t_py_generator::type_to_enum(t_type* type) {
type = get_true_type(type);
Expand All @@ -2829,6 +2945,8 @@ string t_py_generator::type_to_enum(t_type* type) {
return "TType.I64";
case t_base_type::TYPE_DOUBLE:
return "TType.DOUBLE";
case t_base_type::TYPE_UUID:
return "TType.UUID";
default:
throw "compiler error: unhandled type";
}
Expand Down Expand Up @@ -2908,4 +3026,5 @@ THRIFT_REGISTER_GENERATOR(
" Package prefix for generated files.\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"
" type_hints: Generate type hints and type checks in write method, including IntEnum generation.\n"
)

0 comments on commit 819977b

Please sign in to comment.