"""Generates and compiles Python gRPC stubs from proto_library rules.""" load("@rules_proto//proto:defs.bzl", "ProtoInfo") load( "//bazel:protobuf.bzl", "declare_out_files", "get_include_directory", "get_out_dir", "get_plugin_args", "get_proto_arguments", "includes_from_deps", "protos_from_context", ) _GENERATED_PROTO_FORMAT = "{}_pb2.py" _GENERATED_GRPC_PROTO_FORMAT = "{}_pb2_grpc.py" def _generate_py_impl(context): protos = protos_from_context(context) includes = includes_from_deps(context.attr.deps) out_files = declare_out_files(protos, context, _GENERATED_PROTO_FORMAT) tools = [context.executable._protoc] out_dir = get_out_dir(protos, context) arguments = ([ "--python_out={}".format(out_dir.path), ] + [ "--proto_path={}".format(get_include_directory(i)) for i in includes ] + [ "--proto_path={}".format(context.genfiles_dir.path), ]) if context.attr.plugin: arguments += get_plugin_args( context.executable.plugin, [], out_dir.path, False, context.attr.plugin.label.name, ) tools.append(context.executable.plugin) arguments += get_proto_arguments(protos, context.genfiles_dir.path) context.actions.run( inputs = protos + includes, tools = tools, outputs = out_files, executable = context.executable._protoc, arguments = arguments, mnemonic = "ProtocInvocation", ) imports = [] if out_dir.import_path: imports.append("%s/%s/%s" % (context.workspace_name, context.label.package, out_dir.import_path)) return [ DefaultInfo(files = depset(direct = out_files)), PyInfo( transitive_sources = depset(), imports = depset(direct = imports), ), ] _generate_pb2_src = rule( attrs = { "deps": attr.label_list( mandatory = True, allow_empty = False, providers = [ProtoInfo], ), "plugin": attr.label( mandatory = False, executable = True, providers = ["files_to_run"], cfg = "host", ), "_protoc": attr.label( default = Label("//external:protocol_compiler"), providers = ["files_to_run"], executable = True, cfg = "host", ), }, implementation = _generate_py_impl, ) def py_proto_library( name, deps, plugin = None, **kwargs): """Generate python code for a protobuf. Args: name: The name of the target. deps: A list of proto_library dependencies. Must contain a single element. plugin: An optional custom protoc plugin to execute together with generating the protobuf code. **kwargs: Additional arguments to be supplied to the invocation of py_library. """ codegen_target = "_{}_codegen".format(name) if len(deps) != 1: fail("Can only compile a single proto at a time.") _generate_pb2_src( name = codegen_target, deps = deps, plugin = plugin, **kwargs ) native.py_library( name = name, srcs = [":{}".format(codegen_target)], deps = [ "@com_google_protobuf//:protobuf_python", ":{}".format(codegen_target), ], **kwargs ) def _generate_pb2_grpc_src_impl(context): protos = protos_from_context(context) includes = includes_from_deps(context.attr.deps) out_files = declare_out_files(protos, context, _GENERATED_GRPC_PROTO_FORMAT) plugin_flags = ["grpc_2_0"] + context.attr.strip_prefixes arguments = [] tools = [context.executable._protoc, context.executable._grpc_plugin] out_dir = get_out_dir(protos, context) arguments += get_plugin_args( context.executable._grpc_plugin, plugin_flags, out_dir.path, False, ) if context.attr.plugin: arguments += get_plugin_args( context.executable.plugin, [], out_dir.path, False, context.attr.plugin.label.name, ) tools.append(context.executable.plugin) arguments += [ "--proto_path={}".format(get_include_directory(i)) for i in includes ] arguments += ["--proto_path={}".format(context.genfiles_dir.path)] arguments += get_proto_arguments(protos, context.genfiles_dir.path) context.actions.run( inputs = protos + includes, tools = tools, outputs = out_files, executable = context.executable._protoc, arguments = arguments, mnemonic = "ProtocInvocation", ) return [ DefaultInfo(files = depset(direct = out_files)), PyInfo( transitive_sources = depset(), # Imports are already configured by the generated py impl imports = depset(), ), ] _generate_pb2_grpc_src = rule( attrs = { "deps": attr.label_list( mandatory = True, allow_empty = False, providers = [ProtoInfo], ), "strip_prefixes": attr.string_list(), "plugin": attr.label( mandatory = False, executable = True, providers = ["files_to_run"], cfg = "host", ), "_grpc_plugin": attr.label( executable = True, providers = ["files_to_run"], cfg = "host", default = Label("//src/compiler:grpc_python_plugin"), ), "_protoc": attr.label( executable = True, providers = ["files_to_run"], cfg = "host", default = Label("//external:protocol_compiler"), ), }, implementation = _generate_pb2_grpc_src_impl, ) def py_grpc_library( name, srcs, deps, plugin = None, strip_prefixes = [], **kwargs): """Generate python code for gRPC services defined in a protobuf. Args: name: The name of the target. srcs: (List of `labels`) a single proto_library target containing the schema of the service. deps: (List of `labels`) a single py_proto_library target for the proto_library in `srcs`. strip_prefixes: (List of `strings`) If provided, this prefix will be stripped from the beginning of foo_pb2 modules imported by the generated stubs. This is useful in combination with the `imports` attribute of the `py_library` rule. plugin: An optional custom protoc plugin to execute together with generating the gRPC code. **kwargs: Additional arguments to be supplied to the invocation of py_library. """ codegen_grpc_target = "_{}_grpc_codegen".format(name) if len(srcs) != 1: fail("Can only compile a single proto at a time.") if len(deps) != 1: fail("Deps must have length 1.") _generate_pb2_grpc_src( name = codegen_grpc_target, deps = srcs, strip_prefixes = strip_prefixes, plugin = plugin, **kwargs ) native.py_library( name = name, srcs = [ ":{}".format(codegen_grpc_target), ], deps = [ Label("//src/python/grpcio/grpc:grpcio"), ] + deps + [ ":{}".format(codegen_grpc_target), ], **kwargs ) def py2and3_test( name, py_test = native.py_test, **kwargs): """Runs a Python test under both Python 2 and Python 3. Args: name: The name of the test. py_test: The rule to use for each test. **kwargs: Keyword arguments passed directly to the underlying py_test rule. """ if "python_version" in kwargs: fail("Cannot specify 'python_version' in py2and3_test.") names = [name + suffix for suffix in (".python2", ".python3")] python_versions = ["PY2", "PY3"] for case_name, python_version in zip(names, python_versions): py_test( name = case_name, python_version = python_version, **kwargs ) suite_kwargs = {} if "visibility" in kwargs: suite_kwargs["visibility"] = kwargs["visibility"] native.test_suite( name = name, tests = names, **suite_kwargs )