Add new grpc_tools setuptools command BuildPackageProtosStrict

tools/distrib/python/grpcio_tools/grpc_tools/command.py

@@ -15,11 +15,36 @@
 import os
 import pkg_resources
 import sys
+import tempfile
 
 import setuptools
 
 from grpc_tools import protoc
 
+_WELL_KNOWN_PROTOS_INCLUDE = pkg_resources.resource_filename(
+    'grpc_tools', '_proto')
+
+
+def _compile_proto(proto_file,
+                   include='',
+                   python_out='',
+                   grpc_python_out='',
+                   strict=False):
+    command = [
+        'grpc_tools.protoc',
+        '--proto_path={}'.format(include),
+        '--proto_path={}'.format(_WELL_KNOWN_PROTOS_INCLUDE),
+        '--python_out={}'.format(python_out),
+        '--grpc_python_out={}'.format(grpc_python_out),
+    ] + [proto_file]
+    if protoc.main(command) != 0:
+        if strict:
+            sys.stderr.write('error: {} failed'.format(command))
+        else:
+            sys.stderr.write('warning: {} failed'.format(command))
+        return False
+    return True
+
 
 def build_package_protos(package_root):
     proto_files = []
@@ -30,19 +55,49 @@ def build_package_protos(package_root):
                 proto_files.append(
                     os.path.abspath(os.path.join(root, filename)))
 
-    well_known_protos_include = pkg_resources.resource_filename(
-        'grpc_tools', '_proto')
+    for proto_file in proto_files:
+        _compile_proto(
+            proto_file,
+            include=inclusion_root,
+            python_out=inclusion_root,
+            grpc_python_out=inclusion_root,
+            strict=False,
+        )
+
 
+def build_package_protos_strict(package_root):
+    proto_files = []
+    inclusion_root = os.path.abspath(package_root)
+    for root, _, files in os.walk(inclusion_root):
+        for filename in files:
+            if filename.endswith('.proto'):
+                proto_files.append(
+                    os.path.abspath(os.path.join(root, filename)))
+
+    tmp_out_directory = tempfile.mkdtemp()
+    compile_failed = False
     for proto_file in proto_files:
-        command = [
-            'grpc_tools.protoc',
-            '--proto_path={}'.format(inclusion_root),
-            '--proto_path={}'.format(well_known_protos_include),
-            '--python_out={}'.format(inclusion_root),
-            '--grpc_python_out={}'.format(inclusion_root),
-        ] + [proto_file]
-        if protoc.main(command) != 0:
-            raise RuntimeError('error: {} failed'.format(command))
+        # Output all the errors across all the files instead of exiting on the
+        # first error proto file.
+        compile_failed |= not _compile_proto(
+            proto_file,
+            include=inclusion_root,
+            python_out=tmp_out_directory,
+            grpc_python_out=tmp_out_directory,
+            strict=True,
+        )
+
+    if compile_failed:
+        sys.exit(1)
+
+    for proto_file in proto_files:
+        _compile_proto(
+            proto_file,
+            include=inclusion_root,
+            python_out=inclusion_root,
+            grpc_python_out=inclusion_root,
+            strict=False,
+        )
 
 
 class BuildPackageProtos(setuptools.Command):
@@ -63,3 +118,26 @@ class BuildPackageProtos(setuptools.Command):
         # to `self.distribution.package_dir` (and get a key error if it's not
         # there).
         build_package_protos(self.distribution.package_dir[''])
+
+
+class BuildPackageProtosStrict(setuptools.Command):
+    """Command to strictly generate project *_pb2.py modules from proto files.
+
+    The generation will abort if any of the proto files contains error.
+    """
+
+    description = 'strictly build grpc protobuf modules'
+    user_options = []
+
+    def initialize_options(self):
+        pass
+
+    def finalize_options(self):
+        pass
+
+    def run(self):
+        # due to limitations of the proto generator, we require that only *one*
+        # directory is provided as an 'include' directory. We assume it's the '' key
+        # to `self.distribution.package_dir` (and get a key error if it's not
+        # there).
+        build_package_protos_strict(self.distribution.package_dir[''])