From d84a6de3a5c59d049cbdc4ce24ed2a2353512072 Mon Sep 17 00:00:00 2001
From: Richard Belleville <rbellevi@google.com>
Date: Thu, 4 Aug 2022 14:34:00 -0700
Subject: [PATCH] Add type stub generation support to grpcio-tools (#30498)

* Add type stub generation support to grpcio-tools

* Clang format
---
 tools/distrib/python/grpcio_tools/grpc_tools/command.py | 1 +
 tools/distrib/python/grpcio_tools/grpc_tools/main.cc    | 6 ++++++
 2 files changed, 7 insertions(+)

diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/command.py b/tools/distrib/python/grpcio_tools/grpc_tools/command.py
index 4aa22443ae2..3fd16687f4b 100644
--- a/tools/distrib/python/grpcio_tools/grpc_tools/command.py
+++ b/tools/distrib/python/grpcio_tools/grpc_tools/command.py
@@ -38,6 +38,7 @@ def build_package_protos(package_root, strict_mode=False):
             '--proto_path={}'.format(inclusion_root),
             '--proto_path={}'.format(well_known_protos_include),
             '--python_out={}'.format(inclusion_root),
+            '--pyi_out={}'.format(inclusion_root),
             '--grpc_python_out={}'.format(inclusion_root),
         ] + [proto_file]
         if protoc.main(command) != 0:
diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/main.cc b/tools/distrib/python/grpcio_tools/grpc_tools/main.cc
index 2f1c70ffdfa..12b40e8ab4a 100644
--- a/tools/distrib/python/grpcio_tools/grpc_tools/main.cc
+++ b/tools/distrib/python/grpcio_tools/grpc_tools/main.cc
@@ -24,6 +24,7 @@
 #include <google/protobuf/compiler/code_generator.h>
 #include <google/protobuf/compiler/command_line_interface.h>
 #include <google/protobuf/compiler/importer.h>
+#include <google/protobuf/compiler/python/pyi_generator.h>
 #include <google/protobuf/compiler/python/python_generator.h>
 #include <google/protobuf/descriptor.h>
 #include <google/protobuf/io/zero_copy_stream_impl_lite.h>
@@ -49,6 +50,11 @@ int protoc_main(int argc, char* argv[]) {
   cli.RegisterGenerator("--python_out", &py_generator,
                         "Generate Python source file.");
 
+  // pyi files for type checking
+  google::protobuf::compiler::python::PyiGenerator pyi_generator;
+  cli.RegisterGenerator("--pyi_out", &pyi_generator,
+                        "Generate Python pyi stub.");
+
   // gRPC Python
   grpc_python_generator::GeneratorConfiguration grpc_py_config;
   grpc_python_generator::PythonGrpcGenerator grpc_py_generator(grpc_py_config);