mirror of
https://github.com/google/flatbuffers.git
synced 2026-06-04 20:48:59 +00:00
[gRPC] Update the code generator for Python to produce typed handlers (#8326)
* Move `namer.h` and `idl_namer.h` to `include/codegen` so they can be reused from `grpc` dirqectory. * [gRPC] Update the Python generator to produce typed handlers and Python stubs if requested. * [gRPC] Document the newly added compiler flags.
This commit is contained in:
Anton Bobukh
@@ -16,136 +16,365 @@
|
||||
*
|
||||
*/
|
||||
|
||||
#include <map>
|
||||
#include <sstream>
|
||||
|
||||
#include "flatbuffers/util.h"
|
||||
#include "src/compiler/python_generator.h"
|
||||
|
||||
namespace grpc_python_generator {
|
||||
#include <algorithm>
|
||||
#include <cstddef>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
#include "codegen/idl_namer.h"
|
||||
#include "codegen/namer.h"
|
||||
#include "codegen/python.h"
|
||||
#include "flatbuffers/idl.h"
|
||||
#include "flatbuffers/util.h"
|
||||
|
||||
namespace flatbuffers {
|
||||
namespace python {
|
||||
namespace grpc {
|
||||
namespace {
|
||||
|
||||
static grpc::string GenerateMethodType(const grpc_generator::Method *method) {
|
||||
|
||||
if (method->NoStreaming())
|
||||
return "unary_unary";
|
||||
|
||||
if (method->ServerStreaming())
|
||||
return "unary_stream";
|
||||
|
||||
if (method->ClientStreaming())
|
||||
return "stream_unary";
|
||||
|
||||
return "stream_stream";
|
||||
bool ClientStreaming(const RPCCall *method) {
|
||||
const Value *val = method->attributes.Lookup("streaming");
|
||||
return val != nullptr && (val->constant == "client" || val->constant == "bidi");
|
||||
}
|
||||
|
||||
grpc::string GenerateMethodInput(const grpc_generator::Method *method) {
|
||||
|
||||
if (method->NoStreaming() || method->ServerStreaming())
|
||||
return "self, request, context";
|
||||
|
||||
return "self, request_iterator, context";
|
||||
bool ServerStreaming(const RPCCall *method) {
|
||||
const Value *val = method->attributes.Lookup("streaming");
|
||||
return val != nullptr && (val->constant == "server" || val->constant == "bidi");
|
||||
}
|
||||
|
||||
void GenerateStub(const grpc_generator::Service *service,
|
||||
grpc_generator::Printer *printer,
|
||||
std::map<grpc::string, grpc::string> *dictonary) {
|
||||
auto vars = *dictonary;
|
||||
printer->Print(vars, "class $ServiceName$Stub(object):\n");
|
||||
printer->Indent();
|
||||
printer->Print("\"\"\" Interface exported by the server. \"\"\"");
|
||||
printer->Print("\n\n");
|
||||
printer->Print("def __init__(self, channel):\n");
|
||||
printer->Indent();
|
||||
printer->Print("\"\"\" Constructor. \n\n");
|
||||
printer->Print("Args: \n");
|
||||
printer->Print("channel: A grpc.Channel. \n");
|
||||
printer->Print("\"\"\"\n\n");
|
||||
|
||||
for (int j = 0; j < service->method_count(); j++) {
|
||||
auto method = service->method(j);
|
||||
vars["MethodName"] = method->name();
|
||||
vars["MethodType"] = GenerateMethodType(&*method);
|
||||
printer->Print(vars, "self.$MethodName$ = channel.$MethodType$(\n");
|
||||
printer->Indent();
|
||||
printer->Print(vars, "\"/$PATH$$ServiceName$/$MethodName$\"\n");
|
||||
printer->Print(")\n");
|
||||
printer->Outdent();
|
||||
printer->Print("\n");
|
||||
void FormatImports(std::stringstream &ss, const Imports &imports) {
|
||||
std::set<std::string> modules;
|
||||
std::map<std::string, std::set<std::string>> names_by_module;
|
||||
for (const Import &import : imports.imports) {
|
||||
if (import.IsLocal()) continue; // skip all local imports
|
||||
if (import.name == "") {
|
||||
modules.insert(import.module);
|
||||
} else {
|
||||
names_by_module[import.module].insert(import.name);
|
||||
}
|
||||
}
|
||||
printer->Outdent();
|
||||
printer->Outdent();
|
||||
printer->Print("\n");
|
||||
}
|
||||
|
||||
void GenerateServicer(const grpc_generator::Service *service,
|
||||
grpc_generator::Printer *printer,
|
||||
std::map<grpc::string, grpc::string> *dictonary) {
|
||||
auto vars = *dictonary;
|
||||
printer->Print(vars, "class $ServiceName$Servicer(object):\n");
|
||||
printer->Indent();
|
||||
printer->Print("\"\"\" Interface exported by the server. \"\"\"");
|
||||
printer->Print("\n\n");
|
||||
|
||||
for (int j = 0; j < service->method_count(); j++) {
|
||||
auto method = service->method(j);
|
||||
vars["MethodName"] = method->name();
|
||||
vars["MethodInput"] = GenerateMethodInput(&*method);
|
||||
printer->Print(vars, "def $MethodName$($MethodInput$):\n");
|
||||
printer->Indent();
|
||||
printer->Print("context.set_code(grpc.StatusCode.UNIMPLEMENTED)\n");
|
||||
printer->Print("context.set_details('Method not implemented!')\n");
|
||||
printer->Print("raise NotImplementedError('Method not implemented!')\n");
|
||||
printer->Outdent();
|
||||
printer->Print("\n\n");
|
||||
for (const std::string &module : modules) {
|
||||
ss << "import " << module << '\n';
|
||||
}
|
||||
printer->Outdent();
|
||||
printer->Print("\n");
|
||||
|
||||
}
|
||||
|
||||
void GenerateRegister(const grpc_generator::Service *service,
|
||||
grpc_generator::Printer *printer,
|
||||
std::map<grpc::string, grpc::string> *dictonary) {
|
||||
auto vars = *dictonary;
|
||||
printer->Print(vars, "def add_$ServiceName$Servicer_to_server(servicer, server):\n");
|
||||
printer->Indent();
|
||||
printer->Print("rpc_method_handlers = {\n");
|
||||
printer->Indent();
|
||||
for (int j = 0; j < service->method_count(); j++) {
|
||||
auto method = service->method(j);
|
||||
vars["MethodName"] = method->name();
|
||||
vars["MethodType"] = GenerateMethodType(&*method);
|
||||
printer->Print(vars, "'$MethodName$': grpc.$MethodType$_rpc_method_handler(\n");
|
||||
printer->Indent();
|
||||
printer->Print(vars, "servicer.$MethodName$\n");
|
||||
printer->Outdent();
|
||||
printer->Print("),\n");
|
||||
ss << '\n';
|
||||
for (const auto &import : names_by_module) {
|
||||
ss << "from " << import.first << " import ";
|
||||
size_t i = 0;
|
||||
for (const std::string &name : import.second) {
|
||||
if (i > 0) ss << ", ";
|
||||
ss << name;
|
||||
++i;
|
||||
}
|
||||
ss << '\n';
|
||||
}
|
||||
printer->Outdent();
|
||||
printer->Print("}\n");
|
||||
printer->Print(vars, "generic_handler = grpc.method_handlers_generic_handler(\n");
|
||||
printer->Indent();
|
||||
printer->Print(vars, "'$PATH$$ServiceName$', rpc_method_handlers)\n");
|
||||
printer->Outdent();
|
||||
printer->Print("server.add_generic_rpc_handlers((generic_handler,))");
|
||||
printer->Outdent();
|
||||
printer->Print("\n");
|
||||
}
|
||||
} // namespace
|
||||
|
||||
grpc::string Generate(grpc_generator::File *file,
|
||||
const grpc_generator::Service *service) {
|
||||
grpc::string output;
|
||||
std::map<grpc::string, grpc::string> vars;
|
||||
vars["PATH"] = file->package();
|
||||
if (!file->package().empty()) { vars["PATH"].append("."); }
|
||||
vars["ServiceName"] = service->name();
|
||||
auto printer = file->CreatePrinter(&output);
|
||||
GenerateStub(service, &*printer, &vars);
|
||||
GenerateServicer(service, &*printer, &vars);
|
||||
GenerateRegister(service, &*printer, &vars);
|
||||
return output;
|
||||
ss << "\n\n";
|
||||
}
|
||||
|
||||
} // namespace grpc_python_generator
|
||||
bool SaveStub(const std::string &filename, const Imports &imports,
|
||||
const std::string &content) {
|
||||
std::stringstream ss;
|
||||
ss << "# Generated by the gRPC FlatBuffers compiler. DO NOT EDIT!\n"
|
||||
<< '\n'
|
||||
<< "from __future__ import annotations\n"
|
||||
<< '\n';
|
||||
FormatImports(ss, imports);
|
||||
ss << content << '\n';
|
||||
|
||||
EnsureDirExists(StripFileName(filename));
|
||||
return flatbuffers::SaveFile(filename.c_str(), ss.str(), false);
|
||||
}
|
||||
|
||||
bool SaveService(const std::string &filename, const Imports &imports,
|
||||
const std::string &content) {
|
||||
std::stringstream ss;
|
||||
ss << "# Generated by the gRPC FlatBuffers compiler. DO NOT EDIT!\n" << '\n';
|
||||
FormatImports(ss, imports);
|
||||
ss << content << '\n';
|
||||
|
||||
EnsureDirExists(StripFileName(filename));
|
||||
return flatbuffers::SaveFile(filename.c_str(), ss.str(), false);
|
||||
}
|
||||
|
||||
class BaseGenerator {
|
||||
protected:
|
||||
BaseGenerator(const Parser &parser, const Namer::Config &config,
|
||||
const std::string &path, const Version &version)
|
||||
: parser_{parser},
|
||||
namer_{WithFlagOptions(config, parser.opts, path), Keywords(version)},
|
||||
version_{version} {}
|
||||
|
||||
protected:
|
||||
std::string ModuleForFile(const std::string &file) const {
|
||||
std::string module = parser_.opts.include_prefix + StripExtension(file) +
|
||||
parser_.opts.filename_suffix;
|
||||
std::replace(module.begin(), module.end(), '/', '.');
|
||||
return module;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::string ModuleFor(const T *def) const {
|
||||
if (parser_.opts.one_file) return ModuleForFile(def->file);
|
||||
return namer_.NamespacedType(*def);
|
||||
}
|
||||
|
||||
const Parser &parser_;
|
||||
const IdlNamer namer_;
|
||||
const Version version_;
|
||||
};
|
||||
|
||||
class StubGenerator : public BaseGenerator {
|
||||
public:
|
||||
StubGenerator(const Parser &parser, const std::string &path,
|
||||
const Version &version)
|
||||
: BaseGenerator(parser, kStubConfig, path, version) {}
|
||||
|
||||
bool Generate() {
|
||||
Imports imports;
|
||||
std::stringstream stub;
|
||||
for (const ServiceDef *service : parser_.services_.vec) {
|
||||
Generate(stub, service, &imports);
|
||||
}
|
||||
|
||||
std::string filename =
|
||||
namer_.config_.output_path +
|
||||
StripPath(StripExtension(parser_.file_being_parsed_)) + "_grpc" +
|
||||
parser_.opts.grpc_filename_suffix + namer_.config_.filename_extension;
|
||||
|
||||
return SaveStub(filename, imports, stub.str());
|
||||
}
|
||||
|
||||
private:
|
||||
void Generate(std::stringstream &ss, const ServiceDef *service,
|
||||
Imports *imports) {
|
||||
imports->Import("grpc");
|
||||
|
||||
ss << "class " << service->name << "Stub(object):\n"
|
||||
<< " def __init__(self, channel: grpc.Channel) -> None: ...\n";
|
||||
|
||||
for (const RPCCall *method : service->calls.vec) {
|
||||
std::string request = "bytes";
|
||||
std::string response = "bytes";
|
||||
|
||||
if (parser_.opts.grpc_python_typed_handlers) {
|
||||
request = namer_.Type(*method->request);
|
||||
response = namer_.Type(*method->response);
|
||||
|
||||
imports->Import(ModuleFor(method->request), request);
|
||||
imports->Import(ModuleFor(method->response), response);
|
||||
}
|
||||
|
||||
ss << " def " << method->name << "(self, ";
|
||||
if (ClientStreaming(method)) {
|
||||
imports->Import("typing");
|
||||
ss << "request_iterator: typing.Iterator[" << request << "]";
|
||||
} else {
|
||||
ss << "request: " << request;
|
||||
}
|
||||
ss << ") -> ";
|
||||
if (ServerStreaming(method)) {
|
||||
imports->Import("typing");
|
||||
ss << "typing.Iterator[" << response << "]";
|
||||
} else {
|
||||
ss << response;
|
||||
}
|
||||
ss << ": ...\n";
|
||||
}
|
||||
|
||||
ss << "\n\n";
|
||||
ss << "class " << service->name << "Servicer(object):\n";
|
||||
|
||||
for (const RPCCall *method : service->calls.vec) {
|
||||
std::string request = "bytes";
|
||||
std::string response = "bytes";
|
||||
|
||||
if (parser_.opts.grpc_python_typed_handlers) {
|
||||
request = namer_.Type(*method->request);
|
||||
response = namer_.Type(*method->response);
|
||||
|
||||
imports->Import(ModuleFor(method->request), request);
|
||||
imports->Import(ModuleFor(method->response), response);
|
||||
}
|
||||
|
||||
ss << " def " << method->name << "(self, ";
|
||||
if (ClientStreaming(method)) {
|
||||
imports->Import("typing");
|
||||
ss << "request_iterator: typing.Iterator[" << request << "]";
|
||||
} else {
|
||||
ss << "request: " << request;
|
||||
}
|
||||
ss << ", context: grpc.ServicerContext) -> ";
|
||||
if (ServerStreaming(method)) {
|
||||
imports->Import("typing");
|
||||
ss << "typing.Iterator[" << response << "]";
|
||||
} else {
|
||||
ss << response;
|
||||
}
|
||||
ss << ": ...\n";
|
||||
}
|
||||
|
||||
ss << '\n'
|
||||
<< '\n'
|
||||
<< "def add_" << service->name
|
||||
<< "Servicer_to_server(servicer: " << service->name
|
||||
<< "Servicer, server: grpc.Server) -> None: ...\n";
|
||||
}
|
||||
};
|
||||
|
||||
class ServiceGenerator : public BaseGenerator {
|
||||
public:
|
||||
ServiceGenerator(const Parser &parser, const std::string &path,
|
||||
const Version &version)
|
||||
: BaseGenerator(parser, kConfig, path, version) {}
|
||||
|
||||
bool Generate() {
|
||||
Imports imports;
|
||||
std::stringstream ss;
|
||||
|
||||
imports.Import("flatbuffers");
|
||||
|
||||
if (parser_.opts.grpc_python_typed_handlers) {
|
||||
ss << "def _serialize_to_bytes(table):\n"
|
||||
<< " buf = table._tab.Bytes\n"
|
||||
<< " n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, 0)\n"
|
||||
<< " if table._tab.Pos != n:\n"
|
||||
<< " raise ValueError('must be a top-level table')\n"
|
||||
<< " return bytes(buf)\n"
|
||||
<< '\n'
|
||||
<< '\n';
|
||||
}
|
||||
|
||||
for (const ServiceDef *service : parser_.services_.vec) {
|
||||
GenerateStub(ss, service, &imports);
|
||||
GenerateServicer(ss, service, &imports);
|
||||
GenerateRegister(ss, service, &imports);
|
||||
}
|
||||
|
||||
std::string filename =
|
||||
namer_.config_.output_path +
|
||||
StripPath(StripExtension(parser_.file_being_parsed_)) + "_grpc" +
|
||||
parser_.opts.grpc_filename_suffix + namer_.config_.filename_extension;
|
||||
|
||||
return SaveService(filename, imports, ss.str());
|
||||
}
|
||||
|
||||
private:
|
||||
void GenerateStub(std::stringstream &ss, const ServiceDef *service,
|
||||
Imports *imports) {
|
||||
ss << "class " << service->name << "Stub";
|
||||
if (version_.major != 3) ss << "(object)";
|
||||
ss << ":\n"
|
||||
<< " '''Interface exported by the server.'''\n"
|
||||
<< '\n'
|
||||
<< " def __init__(self, channel):\n"
|
||||
<< " '''Constructor.\n"
|
||||
<< '\n'
|
||||
<< " Args:\n"
|
||||
<< " channel: A grpc.Channel.\n"
|
||||
<< " '''\n"
|
||||
<< '\n';
|
||||
|
||||
for (const RPCCall *method : service->calls.vec) {
|
||||
std::string response = namer_.Type(*method->response);
|
||||
|
||||
imports->Import(ModuleFor(method->response), response);
|
||||
|
||||
ss << " self." << method->name << " = channel."
|
||||
<< (ClientStreaming(method) ? "stream" : "unary") << "_"
|
||||
<< (ServerStreaming(method) ? "stream" : "unary") << "(\n"
|
||||
<< " method='/"
|
||||
<< service->defined_namespace->GetFullyQualifiedName(service->name)
|
||||
<< "/" << method->name << "'";
|
||||
|
||||
if (parser_.opts.grpc_python_typed_handlers) {
|
||||
ss << ",\n"
|
||||
<< " request_serializer=_serialize_to_bytes,\n"
|
||||
<< " response_deserializer=" << response << ".GetRootAs";
|
||||
}
|
||||
ss << ")\n\n";
|
||||
}
|
||||
|
||||
ss << '\n';
|
||||
}
|
||||
|
||||
void GenerateServicer(std::stringstream &ss, const ServiceDef *service,
|
||||
Imports *imports) {
|
||||
imports->Import("grpc");
|
||||
|
||||
ss << "class " << service->name << "Servicer";
|
||||
if (version_.major != 3) ss << "(object)";
|
||||
ss << ":\n"
|
||||
<< " '''Interface exported by the server.'''\n"
|
||||
<< '\n';
|
||||
|
||||
for (const RPCCall *method : service->calls.vec) {
|
||||
const std::string request_param =
|
||||
ClientStreaming(method) ? "request_iterator" : "request";
|
||||
ss << " def " << method->name << "(self, " << request_param
|
||||
<< ", context):\n"
|
||||
<< " context.set_code(grpc.StatusCode.UNIMPLEMENTED)\n"
|
||||
<< " context.set_details('Method not implemented!')\n"
|
||||
<< " raise NotImplementedError('Method not implemented!')\n"
|
||||
<< '\n';
|
||||
}
|
||||
|
||||
ss << '\n';
|
||||
}
|
||||
|
||||
void GenerateRegister(std::stringstream &ss, const ServiceDef *service,
|
||||
Imports *imports) {
|
||||
imports->Import("grpc");
|
||||
|
||||
ss << "def add_" << service->name
|
||||
<< "Servicer_to_server(servicer, server):\n"
|
||||
<< " rpc_method_handlers = {\n";
|
||||
|
||||
for (const RPCCall *method : service->calls.vec) {
|
||||
std::string request = namer_.Type(*method->request);
|
||||
|
||||
imports->Import(ModuleFor(method->request), request);
|
||||
|
||||
ss << " '" << method->name << "': grpc."
|
||||
<< (ClientStreaming(method) ? "stream" : "unary") << "_"
|
||||
<< (ServerStreaming(method) ? "stream" : "unary")
|
||||
<< "_rpc_method_handler(\n"
|
||||
<< " servicer." << method->name;
|
||||
|
||||
if (parser_.opts.grpc_python_typed_handlers) {
|
||||
ss << ",\n"
|
||||
<< " request_deserializer=" << request << ".GetRootAs,\n"
|
||||
<< " response_serializer=_serialize_to_bytes";
|
||||
}
|
||||
ss << "),\n";
|
||||
}
|
||||
ss << " }\n"
|
||||
<< '\n'
|
||||
<< " generic_handler = grpc.method_handlers_generic_handler(\n"
|
||||
<< " '"
|
||||
<< service->defined_namespace->GetFullyQualifiedName(service->name)
|
||||
<< "', rpc_method_handlers)\n"
|
||||
<< '\n'
|
||||
<< " server.add_generic_rpc_handlers((generic_handler,))\n"
|
||||
<< '\n';
|
||||
}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
bool Generate(const Parser &parser, const std::string &path,
|
||||
const Version &version) {
|
||||
ServiceGenerator generator{parser, path, version};
|
||||
return generator.Generate();
|
||||
}
|
||||
|
||||
bool GenerateStub(const Parser &parser, const std::string &path,
|
||||
const Version &version) {
|
||||
StubGenerator generator{parser, path, version};
|
||||
return generator.Generate();
|
||||
}
|
||||
|
||||
} // namespace grpc
|
||||
} // namespace python
|
||||
} // namespace flatbuffers
|
||||
|
||||
Reference in New Issue
Block a user