Minimize the change

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

@@ -21,32 +21,8 @@ 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):
+def build_package_protos(package_root, strict_mode=False):
     proto_files = []
     inclusion_root = os.path.abspath(package_root)
     for root, _, files in os.walk(inclusion_root):
@@ -55,82 +31,33 @@ def build_package_protos(package_root):
                 proto_files.append(
                     os.path.abspath(os.path.join(root, filename)))
 
-    for proto_file in proto_files:
-        _compile_proto(
-            proto_file,
-            include=inclusion_root,
-            python_out=inclusion_root,
-            grpc_python_out=inclusion_root,
-            strict=False,
-        )
-
+    well_known_protos_include = pkg_resources.resource_filename(
+        'grpc_tools', '_proto')
 
-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:
-        # 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,
-        )
+        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:
+            if strict_mode:
+                raise Exception('error: {} failed'.format(command))
+            else:
+                sys.stderr.write('warning: {} failed'.format(command))
 
 
 class BuildPackageProtos(setuptools.Command):
     """Command to generate project *_pb2.py modules from proto files."""
 
     description = 'build grpc protobuf modules'
-    user_options = []
+    user_options = [('strict-mode', 's',
+                     'exit with non-zero value if the proto compiling fails.')]
 
     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(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
+        self.strict_mode = False
 
     def finalize_options(self):
         pass
@@ -140,4 +67,7 @@ class BuildPackageProtosStrict(setuptools.Command):
         # 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[''])
+        if self.strict_mode:
+            self.announce('Building Package Protos in Strict Mode')
+        build_package_protos(self.distribution.package_dir[''],
+                             self.strict_mode)