command.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. # Copyright 2016 gRPC authors.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import sys
  16. from grpc_tools import protoc
  17. import pkg_resources
  18. import setuptools
  19. def build_package_protos(package_root, strict_mode=False):
  20. proto_files = []
  21. inclusion_root = os.path.abspath(package_root)
  22. for root, _, files in os.walk(inclusion_root):
  23. for filename in files:
  24. if filename.endswith('.proto'):
  25. proto_files.append(os.path.abspath(os.path.join(root,
  26. filename)))
  27. well_known_protos_include = pkg_resources.resource_filename(
  28. 'grpc_tools', '_proto')
  29. for proto_file in proto_files:
  30. command = [
  31. 'grpc_tools.protoc',
  32. '--proto_path={}'.format(inclusion_root),
  33. '--proto_path={}'.format(well_known_protos_include),
  34. '--python_out={}'.format(inclusion_root),
  35. '--grpc_python_out={}'.format(inclusion_root),
  36. ] + [proto_file]
  37. if protoc.main(command) != 0:
  38. if strict_mode:
  39. raise Exception('error: {} failed'.format(command))
  40. else:
  41. sys.stderr.write('warning: {} failed'.format(command))
  42. class BuildPackageProtos(setuptools.Command):
  43. """Command to generate project *_pb2.py modules from proto files."""
  44. description = 'build grpc protobuf modules'
  45. user_options = [('strict-mode', 's',
  46. 'exit with non-zero value if the proto compiling fails.')]
  47. def initialize_options(self):
  48. self.strict_mode = False
  49. def finalize_options(self):
  50. pass
  51. def run(self):
  52. # due to limitations of the proto generator, we require that only *one*
  53. # directory is provided as an 'include' directory. We assume it's the '' key
  54. # to `self.distribution.package_dir` (and get a key error if it's not
  55. # there).
  56. build_package_protos(self.distribution.package_dir[''],
  57. self.strict_mode)