• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  *
3  * Copyright 2015 gRPC authors.
4  *
5  * Licensed under the Apache License, Version 2.0 (the "License");
6  * you may not use this file except in compliance with the License.
7  * You may obtain a copy of the License at
8  *
9  *     http://www.apache.org/licenses/LICENSE-2.0
10  *
11  * Unless required by applicable law or agreed to in writing, software
12  * distributed under the License is distributed on an "AS IS" BASIS,
13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  * See the License for the specific language governing permissions and
15  * limitations under the License.
16  *
17  */
18 
19 #include <algorithm>
20 #include <cassert>
21 #include <cctype>
22 #include <cstring>
23 #include <fstream>
24 #include <iostream>
25 #include <map>
26 #include <memory>
27 #include <ostream>
28 #include <set>
29 #include <sstream>
30 #include <tuple>
31 #include <vector>
32 
33 #include "src/compiler/config.h"
34 #include "src/compiler/generator_helpers.h"
35 #include "src/compiler/protobuf_plugin.h"
36 #include "src/compiler/python_generator.h"
37 #include "src/compiler/python_generator_helpers.h"
38 #include "src/compiler/python_private_generator.h"
39 
40 using grpc::protobuf::FileDescriptor;
41 using grpc::protobuf::compiler::GeneratorContext;
42 using grpc::protobuf::io::CodedOutputStream;
43 using grpc::protobuf::io::ZeroCopyOutputStream;
44 using std::make_pair;
45 using std::map;
46 using std::pair;
47 using std::replace;
48 using std::set;
49 using std::tuple;
50 using std::vector;
51 
52 namespace grpc_python_generator {
53 
54 grpc::string generator_file_name;
55 
56 namespace {
57 
58 typedef map<grpc::string, grpc::string> StringMap;
59 typedef vector<grpc::string> StringVector;
60 typedef tuple<grpc::string, grpc::string> StringPair;
61 typedef set<StringPair> StringPairSet;
62 
63 // Provides RAII indentation handling. Use as:
64 // {
65 //   IndentScope raii_my_indent_var_name_here(my_py_printer);
66 //   // constructor indented my_py_printer
67 //   ...
68 //   // destructor called at end of scope, un-indenting my_py_printer
69 // }
70 class IndentScope {
71  public:
IndentScope(grpc_generator::Printer * printer)72   explicit IndentScope(grpc_generator::Printer* printer) : printer_(printer) {
73     printer_->Indent();
74   }
75 
~IndentScope()76   ~IndentScope() { printer_->Outdent(); }
77 
78  private:
79   grpc_generator::Printer* printer_;
80 };
81 
PrivateGenerator(const GeneratorConfiguration & config,const grpc_generator::File * file)82 PrivateGenerator::PrivateGenerator(const GeneratorConfiguration& config,
83                                    const grpc_generator::File* file)
84     : config(config), file(file) {}
85 
PrintAllComments(StringVector comments,grpc_generator::Printer * out)86 void PrivateGenerator::PrintAllComments(StringVector comments,
87                                         grpc_generator::Printer* out) {
88   if (comments.empty()) {
89     // Python requires code structures like class and def to have
90     // a body, even if it is just "pass" or a docstring.  We need
91     // to ensure not to generate empty bodies. We could do something
92     // smarter and more sophisticated, but at the moment, if there is
93     // no docstring to print, we simply emit "pass" to ensure validity
94     // of the generated code.
95     out->Print("# missing associated documentation comment in .proto file\n");
96     out->Print("pass\n");
97     return;
98   }
99   out->Print("\"\"\"");
100   for (StringVector::iterator it = comments.begin(); it != comments.end();
101        ++it) {
102     size_t start_pos = it->find_first_not_of(' ');
103     if (start_pos != grpc::string::npos) {
104       out->PrintRaw(it->c_str() + start_pos);
105     }
106     out->Print("\n");
107   }
108   out->Print("\"\"\"\n");
109 }
110 
PrintBetaServicer(const grpc_generator::Service * service,grpc_generator::Printer * out)111 bool PrivateGenerator::PrintBetaServicer(const grpc_generator::Service* service,
112                                          grpc_generator::Printer* out) {
113   StringMap service_dict;
114   service_dict["Service"] = service->name();
115   out->Print("\n\n");
116   out->Print(service_dict, "class Beta$Service$Servicer(object):\n");
117   {
118     IndentScope raii_class_indent(out);
119     out->Print(
120         "\"\"\"The Beta API is deprecated for 0.15.0 and later.\n"
121         "\nIt is recommended to use the GA API (classes and functions in this\n"
122         "file not marked beta) for all further purposes. This class was "
123         "generated\n"
124         "only to ease transition from grpcio<0.15.0 to "
125         "grpcio>=0.15.0.\"\"\"\n");
126     StringVector service_comments = service->GetAllComments();
127     PrintAllComments(service_comments, out);
128     for (int i = 0; i < service->method_count(); ++i) {
129       auto method = service->method(i);
130       grpc::string arg_name =
131           method->ClientStreaming() ? "request_iterator" : "request";
132       StringMap method_dict;
133       method_dict["Method"] = method->name();
134       method_dict["ArgName"] = arg_name;
135       out->Print(method_dict, "def $Method$(self, $ArgName$, context):\n");
136       {
137         IndentScope raii_method_indent(out);
138         StringVector method_comments = method->GetAllComments();
139         PrintAllComments(method_comments, out);
140         out->Print("context.code(beta_interfaces.StatusCode.UNIMPLEMENTED)\n");
141       }
142     }
143   }
144   return true;
145 }
146 
PrintBetaStub(const grpc_generator::Service * service,grpc_generator::Printer * out)147 bool PrivateGenerator::PrintBetaStub(const grpc_generator::Service* service,
148                                      grpc_generator::Printer* out) {
149   StringMap service_dict;
150   service_dict["Service"] = service->name();
151   out->Print("\n\n");
152   out->Print(service_dict, "class Beta$Service$Stub(object):\n");
153   {
154     IndentScope raii_class_indent(out);
155     out->Print(
156         "\"\"\"The Beta API is deprecated for 0.15.0 and later.\n"
157         "\nIt is recommended to use the GA API (classes and functions in this\n"
158         "file not marked beta) for all further purposes. This class was "
159         "generated\n"
160         "only to ease transition from grpcio<0.15.0 to "
161         "grpcio>=0.15.0.\"\"\"\n");
162     StringVector service_comments = service->GetAllComments();
163     PrintAllComments(service_comments, out);
164     for (int i = 0; i < service->method_count(); ++i) {
165       auto method = service->method(i);
166       grpc::string arg_name =
167           method->ClientStreaming() ? "request_iterator" : "request";
168       StringMap method_dict;
169       method_dict["Method"] = method->name();
170       method_dict["ArgName"] = arg_name;
171       out->Print(method_dict,
172                  "def $Method$(self, $ArgName$, timeout, metadata=None, "
173                  "with_call=False, protocol_options=None):\n");
174       {
175         IndentScope raii_method_indent(out);
176         StringVector method_comments = method->GetAllComments();
177         PrintAllComments(method_comments, out);
178         out->Print("raise NotImplementedError()\n");
179       }
180       if (!method->ServerStreaming()) {
181         out->Print(method_dict, "$Method$.future = None\n");
182       }
183     }
184   }
185   return true;
186 }
187 
PrintBetaServerFactory(const grpc::string & package_qualified_service_name,const grpc_generator::Service * service,grpc_generator::Printer * out)188 bool PrivateGenerator::PrintBetaServerFactory(
189     const grpc::string& package_qualified_service_name,
190     const grpc_generator::Service* service, grpc_generator::Printer* out) {
191   StringMap service_dict;
192   service_dict["Service"] = service->name();
193   out->Print("\n\n");
194   out->Print(service_dict,
195              "def beta_create_$Service$_server(servicer, pool=None, "
196              "pool_size=None, default_timeout=None, maximum_timeout=None):\n");
197   {
198     IndentScope raii_create_server_indent(out);
199     out->Print(
200         "\"\"\"The Beta API is deprecated for 0.15.0 and later.\n"
201         "\nIt is recommended to use the GA API (classes and functions in this\n"
202         "file not marked beta) for all further purposes. This function was\n"
203         "generated only to ease transition from grpcio<0.15.0 to grpcio>=0.15.0"
204         "\"\"\"\n");
205     StringMap method_implementation_constructors;
206     StringMap input_message_modules_and_classes;
207     StringMap output_message_modules_and_classes;
208     for (int i = 0; i < service->method_count(); ++i) {
209       auto method = service->method(i);
210       const grpc::string method_implementation_constructor =
211           grpc::string(method->ClientStreaming() ? "stream_" : "unary_") +
212           grpc::string(method->ServerStreaming() ? "stream_" : "unary_") +
213           "inline";
214       grpc::string input_message_module_and_class;
215       if (!method->get_module_and_message_path_input(
216               &input_message_module_and_class, generator_file_name,
217               generate_in_pb2_grpc, config.import_prefix)) {
218         return false;
219       }
220       grpc::string output_message_module_and_class;
221       if (!method->get_module_and_message_path_output(
222               &output_message_module_and_class, generator_file_name,
223               generate_in_pb2_grpc, config.import_prefix)) {
224         return false;
225       }
226       method_implementation_constructors.insert(
227           make_pair(method->name(), method_implementation_constructor));
228       input_message_modules_and_classes.insert(
229           make_pair(method->name(), input_message_module_and_class));
230       output_message_modules_and_classes.insert(
231           make_pair(method->name(), output_message_module_and_class));
232     }
233     StringMap method_dict;
234     method_dict["PackageQualifiedServiceName"] = package_qualified_service_name;
235     out->Print("request_deserializers = {\n");
236     for (StringMap::iterator name_and_input_module_class_pair =
237              input_message_modules_and_classes.begin();
238          name_and_input_module_class_pair !=
239          input_message_modules_and_classes.end();
240          name_and_input_module_class_pair++) {
241       method_dict["MethodName"] = name_and_input_module_class_pair->first;
242       method_dict["InputTypeModuleAndClass"] =
243           name_and_input_module_class_pair->second;
244       IndentScope raii_indent(out);
245       out->Print(method_dict,
246                  "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
247                  "$InputTypeModuleAndClass$.FromString,\n");
248     }
249     out->Print("}\n");
250     out->Print("response_serializers = {\n");
251     for (StringMap::iterator name_and_output_module_class_pair =
252              output_message_modules_and_classes.begin();
253          name_and_output_module_class_pair !=
254          output_message_modules_and_classes.end();
255          name_and_output_module_class_pair++) {
256       method_dict["MethodName"] = name_and_output_module_class_pair->first;
257       method_dict["OutputTypeModuleAndClass"] =
258           name_and_output_module_class_pair->second;
259       IndentScope raii_indent(out);
260       out->Print(method_dict,
261                  "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
262                  "$OutputTypeModuleAndClass$.SerializeToString,\n");
263     }
264     out->Print("}\n");
265     out->Print("method_implementations = {\n");
266     for (StringMap::iterator name_and_implementation_constructor =
267              method_implementation_constructors.begin();
268          name_and_implementation_constructor !=
269          method_implementation_constructors.end();
270          name_and_implementation_constructor++) {
271       method_dict["Method"] = name_and_implementation_constructor->first;
272       method_dict["Constructor"] = name_and_implementation_constructor->second;
273       IndentScope raii_descriptions_indent(out);
274       const grpc::string method_name =
275           name_and_implementation_constructor->first;
276       out->Print(method_dict,
277                  "(\'$PackageQualifiedServiceName$\', \'$Method$\'): "
278                  "face_utilities.$Constructor$(servicer.$Method$),\n");
279     }
280     out->Print("}\n");
281     out->Print(
282         "server_options = beta_implementations.server_options("
283         "request_deserializers=request_deserializers, "
284         "response_serializers=response_serializers, "
285         "thread_pool=pool, thread_pool_size=pool_size, "
286         "default_timeout=default_timeout, "
287         "maximum_timeout=maximum_timeout)\n");
288     out->Print(
289         "return beta_implementations.server(method_implementations, "
290         "options=server_options)\n");
291   }
292   return true;
293 }
294 
PrintBetaStubFactory(const grpc::string & package_qualified_service_name,const grpc_generator::Service * service,grpc_generator::Printer * out)295 bool PrivateGenerator::PrintBetaStubFactory(
296     const grpc::string& package_qualified_service_name,
297     const grpc_generator::Service* service, grpc_generator::Printer* out) {
298   StringMap dict;
299   dict["Service"] = service->name();
300   out->Print("\n\n");
301   out->Print(dict,
302              "def beta_create_$Service$_stub(channel, host=None,"
303              " metadata_transformer=None, pool=None, pool_size=None):\n");
304   {
305     IndentScope raii_create_server_indent(out);
306     out->Print(
307         "\"\"\"The Beta API is deprecated for 0.15.0 and later.\n"
308         "\nIt is recommended to use the GA API (classes and functions in this\n"
309         "file not marked beta) for all further purposes. This function was\n"
310         "generated only to ease transition from grpcio<0.15.0 to grpcio>=0.15.0"
311         "\"\"\"\n");
312     StringMap method_cardinalities;
313     StringMap input_message_modules_and_classes;
314     StringMap output_message_modules_and_classes;
315     for (int i = 0; i < service->method_count(); ++i) {
316       auto method = service->method(i);
317       const grpc::string method_cardinality =
318           grpc::string(method->ClientStreaming() ? "STREAM" : "UNARY") + "_" +
319           grpc::string(method->ServerStreaming() ? "STREAM" : "UNARY");
320       grpc::string input_message_module_and_class;
321       if (!method->get_module_and_message_path_input(
322               &input_message_module_and_class, generator_file_name,
323               generate_in_pb2_grpc, config.import_prefix)) {
324         return false;
325       }
326       grpc::string output_message_module_and_class;
327       if (!method->get_module_and_message_path_output(
328               &output_message_module_and_class, generator_file_name,
329               generate_in_pb2_grpc, config.import_prefix)) {
330         return false;
331       }
332       method_cardinalities.insert(
333           make_pair(method->name(), method_cardinality));
334       input_message_modules_and_classes.insert(
335           make_pair(method->name(), input_message_module_and_class));
336       output_message_modules_and_classes.insert(
337           make_pair(method->name(), output_message_module_and_class));
338     }
339     StringMap method_dict;
340     method_dict["PackageQualifiedServiceName"] = package_qualified_service_name;
341     out->Print("request_serializers = {\n");
342     for (StringMap::iterator name_and_input_module_class_pair =
343              input_message_modules_and_classes.begin();
344          name_and_input_module_class_pair !=
345          input_message_modules_and_classes.end();
346          name_and_input_module_class_pair++) {
347       method_dict["MethodName"] = name_and_input_module_class_pair->first;
348       method_dict["InputTypeModuleAndClass"] =
349           name_and_input_module_class_pair->second;
350       IndentScope raii_indent(out);
351       out->Print(method_dict,
352                  "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
353                  "$InputTypeModuleAndClass$.SerializeToString,\n");
354     }
355     out->Print("}\n");
356     out->Print("response_deserializers = {\n");
357     for (StringMap::iterator name_and_output_module_class_pair =
358              output_message_modules_and_classes.begin();
359          name_and_output_module_class_pair !=
360          output_message_modules_and_classes.end();
361          name_and_output_module_class_pair++) {
362       method_dict["MethodName"] = name_and_output_module_class_pair->first;
363       method_dict["OutputTypeModuleAndClass"] =
364           name_and_output_module_class_pair->second;
365       IndentScope raii_indent(out);
366       out->Print(method_dict,
367                  "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
368                  "$OutputTypeModuleAndClass$.FromString,\n");
369     }
370     out->Print("}\n");
371     out->Print("cardinalities = {\n");
372     for (StringMap::iterator name_and_cardinality =
373              method_cardinalities.begin();
374          name_and_cardinality != method_cardinalities.end();
375          name_and_cardinality++) {
376       method_dict["Method"] = name_and_cardinality->first;
377       method_dict["Cardinality"] = name_and_cardinality->second;
378       IndentScope raii_descriptions_indent(out);
379       out->Print(method_dict,
380                  "\'$Method$\': cardinality.Cardinality.$Cardinality$,\n");
381     }
382     out->Print("}\n");
383     out->Print(
384         "stub_options = beta_implementations.stub_options("
385         "host=host, metadata_transformer=metadata_transformer, "
386         "request_serializers=request_serializers, "
387         "response_deserializers=response_deserializers, "
388         "thread_pool=pool, thread_pool_size=pool_size)\n");
389     out->Print(method_dict,
390                "return beta_implementations.dynamic_stub(channel, "
391                "\'$PackageQualifiedServiceName$\', "
392                "cardinalities, options=stub_options)\n");
393   }
394   return true;
395 }
396 
PrintStub(const grpc::string & package_qualified_service_name,const grpc_generator::Service * service,grpc_generator::Printer * out)397 bool PrivateGenerator::PrintStub(
398     const grpc::string& package_qualified_service_name,
399     const grpc_generator::Service* service, grpc_generator::Printer* out) {
400   StringMap dict;
401   dict["Service"] = service->name();
402   out->Print("\n\n");
403   out->Print(dict, "class $Service$Stub(object):\n");
404   {
405     IndentScope raii_class_indent(out);
406     StringVector service_comments = service->GetAllComments();
407     PrintAllComments(service_comments, out);
408     out->Print("\n");
409     out->Print("def __init__(self, channel):\n");
410     {
411       IndentScope raii_init_indent(out);
412       out->Print("\"\"\"Constructor.\n");
413       out->Print("\n");
414       out->Print("Args:\n");
415       {
416         IndentScope raii_args_indent(out);
417         out->Print("channel: A grpc.Channel.\n");
418       }
419       out->Print("\"\"\"\n");
420       for (int i = 0; i < service->method_count(); ++i) {
421         auto method = service->method(i);
422         grpc::string multi_callable_constructor =
423             grpc::string(method->ClientStreaming() ? "stream" : "unary") + "_" +
424             grpc::string(method->ServerStreaming() ? "stream" : "unary");
425         grpc::string request_module_and_class;
426         if (!method->get_module_and_message_path_input(
427                 &request_module_and_class, generator_file_name,
428                 generate_in_pb2_grpc, config.import_prefix)) {
429           return false;
430         }
431         grpc::string response_module_and_class;
432         if (!method->get_module_and_message_path_output(
433                 &response_module_and_class, generator_file_name,
434                 generate_in_pb2_grpc, config.import_prefix)) {
435           return false;
436         }
437         StringMap method_dict;
438         method_dict["Method"] = method->name();
439         method_dict["MultiCallableConstructor"] = multi_callable_constructor;
440         out->Print(method_dict,
441                    "self.$Method$ = channel.$MultiCallableConstructor$(\n");
442         {
443           method_dict["PackageQualifiedService"] =
444               package_qualified_service_name;
445           method_dict["RequestModuleAndClass"] = request_module_and_class;
446           method_dict["ResponseModuleAndClass"] = response_module_and_class;
447           IndentScope raii_first_attribute_indent(out);
448           IndentScope raii_second_attribute_indent(out);
449           out->Print(method_dict, "'/$PackageQualifiedService$/$Method$',\n");
450           out->Print(method_dict,
451                      "request_serializer=$RequestModuleAndClass$."
452                      "SerializeToString,\n");
453           out->Print(
454               method_dict,
455               "response_deserializer=$ResponseModuleAndClass$.FromString,\n");
456           out->Print(")\n");
457         }
458       }
459     }
460   }
461   return true;
462 }
463 
PrintServicer(const grpc_generator::Service * service,grpc_generator::Printer * out)464 bool PrivateGenerator::PrintServicer(const grpc_generator::Service* service,
465                                      grpc_generator::Printer* out) {
466   StringMap service_dict;
467   service_dict["Service"] = service->name();
468   out->Print("\n\n");
469   out->Print(service_dict, "class $Service$Servicer(object):\n");
470   {
471     IndentScope raii_class_indent(out);
472     StringVector service_comments = service->GetAllComments();
473     PrintAllComments(service_comments, out);
474     for (int i = 0; i < service->method_count(); ++i) {
475       auto method = service->method(i);
476       grpc::string arg_name =
477           method->ClientStreaming() ? "request_iterator" : "request";
478       StringMap method_dict;
479       method_dict["Method"] = method->name();
480       method_dict["ArgName"] = arg_name;
481       out->Print("\n");
482       out->Print(method_dict, "def $Method$(self, $ArgName$, context):\n");
483       {
484         IndentScope raii_method_indent(out);
485         StringVector method_comments = method->GetAllComments();
486         PrintAllComments(method_comments, out);
487         out->Print("context.set_code(grpc.StatusCode.UNIMPLEMENTED)\n");
488         out->Print("context.set_details('Method not implemented!')\n");
489         out->Print("raise NotImplementedError('Method not implemented!')\n");
490       }
491     }
492   }
493   return true;
494 }
495 
PrintAddServicerToServer(const grpc::string & package_qualified_service_name,const grpc_generator::Service * service,grpc_generator::Printer * out)496 bool PrivateGenerator::PrintAddServicerToServer(
497     const grpc::string& package_qualified_service_name,
498     const grpc_generator::Service* service, grpc_generator::Printer* out) {
499   StringMap service_dict;
500   service_dict["Service"] = service->name();
501   out->Print("\n\n");
502   out->Print(service_dict,
503              "def add_$Service$Servicer_to_server(servicer, server):\n");
504   {
505     IndentScope raii_class_indent(out);
506     out->Print("rpc_method_handlers = {\n");
507     {
508       IndentScope raii_dict_first_indent(out);
509       IndentScope raii_dict_second_indent(out);
510       for (int i = 0; i < service->method_count(); ++i) {
511         auto method = service->method(i);
512         grpc::string method_handler_constructor =
513             grpc::string(method->ClientStreaming() ? "stream" : "unary") + "_" +
514             grpc::string(method->ServerStreaming() ? "stream" : "unary") +
515             "_rpc_method_handler";
516         grpc::string request_module_and_class;
517         if (!method->get_module_and_message_path_input(
518                 &request_module_and_class, generator_file_name,
519                 generate_in_pb2_grpc, config.import_prefix)) {
520           return false;
521         }
522         grpc::string response_module_and_class;
523         if (!method->get_module_and_message_path_output(
524                 &response_module_and_class, generator_file_name,
525                 generate_in_pb2_grpc, config.import_prefix)) {
526           return false;
527         }
528         StringMap method_dict;
529         method_dict["Method"] = method->name();
530         method_dict["MethodHandlerConstructor"] = method_handler_constructor;
531         method_dict["RequestModuleAndClass"] = request_module_and_class;
532         method_dict["ResponseModuleAndClass"] = response_module_and_class;
533         out->Print(method_dict,
534                    "'$Method$': grpc.$MethodHandlerConstructor$(\n");
535         {
536           IndentScope raii_call_first_indent(out);
537           IndentScope raii_call_second_indent(out);
538           out->Print(method_dict, "servicer.$Method$,\n");
539           out->Print(
540               method_dict,
541               "request_deserializer=$RequestModuleAndClass$.FromString,\n");
542           out->Print(
543               method_dict,
544               "response_serializer=$ResponseModuleAndClass$.SerializeToString,"
545               "\n");
546         }
547         out->Print("),\n");
548       }
549     }
550     StringMap method_dict;
551     method_dict["PackageQualifiedServiceName"] = package_qualified_service_name;
552     out->Print("}\n");
553     out->Print("generic_handler = grpc.method_handlers_generic_handler(\n");
554     {
555       IndentScope raii_call_first_indent(out);
556       IndentScope raii_call_second_indent(out);
557       out->Print(method_dict,
558                  "'$PackageQualifiedServiceName$', rpc_method_handlers)\n");
559     }
560     out->Print("server.add_generic_rpc_handlers((generic_handler,))\n");
561   }
562   return true;
563 }
564 
PrintBetaPreamble(grpc_generator::Printer * out)565 bool PrivateGenerator::PrintBetaPreamble(grpc_generator::Printer* out) {
566   StringMap var;
567   var["Package"] = config.beta_package_root;
568   out->Print(var,
569              "from $Package$ import implementations as beta_implementations\n");
570   out->Print(var, "from $Package$ import interfaces as beta_interfaces\n");
571   out->Print("from grpc.framework.common import cardinality\n");
572   out->Print(
573       "from grpc.framework.interfaces.face import utilities as "
574       "face_utilities\n");
575   return true;
576 }
577 
PrintPreamble(grpc_generator::Printer * out)578 bool PrivateGenerator::PrintPreamble(grpc_generator::Printer* out) {
579   StringMap var;
580   var["Package"] = config.grpc_package_root;
581   out->Print(var, "import $Package$\n");
582   if (generate_in_pb2_grpc) {
583     out->Print("\n");
584     StringPairSet imports_set;
585     for (int i = 0; i < file->service_count(); ++i) {
586       auto service = file->service(i);
587       for (int j = 0; j < service->method_count(); ++j) {
588         auto method = service.get()->method(j);
589 
590         grpc::string input_type_file_name = method->get_input_type_name();
591         grpc::string input_module_name =
592             ModuleName(input_type_file_name, config.import_prefix);
593         grpc::string input_module_alias =
594             ModuleAlias(input_type_file_name, config.import_prefix);
595         imports_set.insert(
596             std::make_tuple(input_module_name, input_module_alias));
597 
598         grpc::string output_type_file_name = method->get_output_type_name();
599         grpc::string output_module_name =
600             ModuleName(output_type_file_name, config.import_prefix);
601         grpc::string output_module_alias =
602             ModuleAlias(output_type_file_name, config.import_prefix);
603         imports_set.insert(
604             std::make_tuple(output_module_name, output_module_alias));
605       }
606     }
607 
608     for (StringPairSet::iterator it = imports_set.begin();
609          it != imports_set.end(); ++it) {
610       auto module_name = std::get<0>(*it);
611       var["ModuleAlias"] = std::get<1>(*it);
612       const size_t last_dot_pos = module_name.rfind('.');
613       if (last_dot_pos == grpc::string::npos) {
614         var["ImportStatement"] = "import " + module_name;
615       } else {
616         var["ImportStatement"] = "from " + module_name.substr(0, last_dot_pos) +
617                                  " import " +
618                                  module_name.substr(last_dot_pos + 1);
619       }
620       out->Print(var, "$ImportStatement$ as $ModuleAlias$\n");
621     }
622   }
623   return true;
624 }
625 
PrintGAServices(grpc_generator::Printer * out)626 bool PrivateGenerator::PrintGAServices(grpc_generator::Printer* out) {
627   grpc::string package = file->package();
628   if (!package.empty()) {
629     package = package.append(".");
630   }
631   for (int i = 0; i < file->service_count(); ++i) {
632     auto service = file->service(i);
633     grpc::string package_qualified_service_name = package + service->name();
634     if (!(PrintStub(package_qualified_service_name, service.get(), out) &&
635           PrintServicer(service.get(), out) &&
636           PrintAddServicerToServer(package_qualified_service_name,
637                                    service.get(), out))) {
638       return false;
639     }
640   }
641   return true;
642 }
643 
PrintBetaServices(grpc_generator::Printer * out)644 bool PrivateGenerator::PrintBetaServices(grpc_generator::Printer* out) {
645   grpc::string package = file->package();
646   if (!package.empty()) {
647     package = package.append(".");
648   }
649   for (int i = 0; i < file->service_count(); ++i) {
650     auto service = file->service(i);
651     grpc::string package_qualified_service_name = package + service->name();
652     if (!(PrintBetaServicer(service.get(), out) &&
653           PrintBetaStub(service.get(), out) &&
654           PrintBetaServerFactory(package_qualified_service_name, service.get(),
655                                  out) &&
656           PrintBetaStubFactory(package_qualified_service_name, service.get(),
657                                out))) {
658       return false;
659     }
660   }
661   return true;
662 }
663 
GetGrpcServices()664 pair<bool, grpc::string> PrivateGenerator::GetGrpcServices() {
665   grpc::string output;
666   {
667     // Scope the output stream so it closes and finalizes output to the string.
668     auto out = file->CreatePrinter(&output);
669     if (generate_in_pb2_grpc) {
670       out->Print(
671           "# Generated by the gRPC Python protocol compiler plugin. "
672           "DO NOT EDIT!\n");
673       if (!PrintPreamble(out.get())) {
674         return make_pair(false, "");
675       }
676       if (!PrintGAServices(out.get())) {
677         return make_pair(false, "");
678       }
679     } else {
680       out->Print("try:\n");
681       {
682         IndentScope raii_dict_try_indent(out.get());
683         out->Print(
684             "# THESE ELEMENTS WILL BE DEPRECATED.\n"
685             "# Please use the generated *_pb2_grpc.py files instead.\n");
686         if (!PrintPreamble(out.get())) {
687           return make_pair(false, "");
688         }
689         if (!PrintBetaPreamble(out.get())) {
690           return make_pair(false, "");
691         }
692         if (!PrintGAServices(out.get())) {
693           return make_pair(false, "");
694         }
695         if (!PrintBetaServices(out.get())) {
696           return make_pair(false, "");
697         }
698       }
699       out->Print("except ImportError:\n");
700       {
701         IndentScope raii_dict_except_indent(out.get());
702         out->Print("pass");
703       }
704     }
705   }
706   return make_pair(true, std::move(output));
707 }
708 
709 }  // namespace
710 
GeneratorConfiguration()711 GeneratorConfiguration::GeneratorConfiguration()
712     : grpc_package_root("grpc"),
713       beta_package_root("grpc.beta"),
714       import_prefix("") {}
715 
PythonGrpcGenerator(const GeneratorConfiguration & config)716 PythonGrpcGenerator::PythonGrpcGenerator(const GeneratorConfiguration& config)
717     : config_(config) {}
718 
~PythonGrpcGenerator()719 PythonGrpcGenerator::~PythonGrpcGenerator() {}
720 
GenerateGrpc(GeneratorContext * context,PrivateGenerator & generator,grpc::string file_name,bool generate_in_pb2_grpc)721 static bool GenerateGrpc(GeneratorContext* context, PrivateGenerator& generator,
722                          grpc::string file_name, bool generate_in_pb2_grpc) {
723   bool success;
724   std::unique_ptr<ZeroCopyOutputStream> output;
725   std::unique_ptr<CodedOutputStream> coded_output;
726   grpc::string grpc_code;
727 
728   if (generate_in_pb2_grpc) {
729     output.reset(context->Open(file_name));
730     generator.generate_in_pb2_grpc = true;
731   } else {
732     output.reset(context->OpenForInsert(file_name, "module_scope"));
733     generator.generate_in_pb2_grpc = false;
734   }
735 
736   coded_output.reset(new CodedOutputStream(output.get()));
737   tie(success, grpc_code) = generator.GetGrpcServices();
738 
739   if (success) {
740     coded_output->WriteRaw(grpc_code.data(), grpc_code.size());
741     return true;
742   } else {
743     return false;
744   }
745 }
746 
Generate(const FileDescriptor * file,const grpc::string & parameter,GeneratorContext * context,grpc::string * error) const747 bool PythonGrpcGenerator::Generate(const FileDescriptor* file,
748                                    const grpc::string& parameter,
749                                    GeneratorContext* context,
750                                    grpc::string* error) const {
751   // Get output file name.
752   grpc::string pb2_file_name;
753   grpc::string pb2_grpc_file_name;
754   static const int proto_suffix_length = strlen(".proto");
755   if (file->name().size() > static_cast<size_t>(proto_suffix_length) &&
756       file->name().find_last_of(".proto") == file->name().size() - 1) {
757     grpc::string base =
758         file->name().substr(0, file->name().size() - proto_suffix_length);
759     std::replace(base.begin(), base.end(), '-', '_');
760     pb2_file_name = base + "_pb2.py";
761     pb2_grpc_file_name = base + "_pb2_grpc.py";
762   } else {
763     *error = "Invalid proto file name. Proto file must end with .proto";
764     return false;
765   }
766   generator_file_name = file->name();
767 
768   ProtoBufFile pbfile(file);
769   PrivateGenerator generator(config_, &pbfile);
770   if (parameter == "" || parameter == "grpc_2_0") {
771     return GenerateGrpc(context, generator, pb2_grpc_file_name, true);
772   } else if (parameter == "grpc_1_0") {
773     return GenerateGrpc(context, generator, pb2_grpc_file_name, true) &&
774            GenerateGrpc(context, generator, pb2_file_name, false);
775   } else {
776     *error = "Invalid parameter '" + parameter + "'.";
777     return false;
778   }
779 }
780 
781 }  // namespace grpc_python_generator
782