From 2dbabc7a8f92d5fe6d3aef8a30297979e5073a9b Mon Sep 17 00:00:00 2001 From: Lidi Zheng <lidiz@google.com> Date: Wed, 16 Jan 2019 12:07:50 -0800 Subject: [PATCH] Add new grpc_tools setuptools command BuildPackageProtosStrict --- .../python/grpcio_tools/grpc_tools/command.py | 100 ++++++++++++++++-- 1 file changed, 89 insertions(+), 11 deletions(-) diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/command.py b/tools/distrib/python/grpcio_tools/grpc_tools/command.py index 93503c4cc36..ee311144225 100644 --- a/tools/distrib/python/grpcio_tools/grpc_tools/command.py +++ b/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[''])