1"""Generates and compiles Python gRPC stubs from proto_library rules.""" 2 3load("@rules_proto//proto:defs.bzl", "ProtoInfo") 4load( 5 "//bazel:protobuf.bzl", 6 "declare_out_files", 7 "get_include_directory", 8 "get_out_dir", 9 "get_plugin_args", 10 "get_proto_arguments", 11 "includes_from_deps", 12 "protos_from_context", 13) 14 15_GENERATED_PROTO_FORMAT = "{}_pb2.py" 16_GENERATED_GRPC_PROTO_FORMAT = "{}_pb2_grpc.py" 17 18def _generate_py_impl(context): 19 protos = protos_from_context(context) 20 includes = includes_from_deps(context.attr.deps) 21 out_files = declare_out_files(protos, context, _GENERATED_PROTO_FORMAT) 22 tools = [context.executable._protoc] 23 24 out_dir = get_out_dir(protos, context) 25 arguments = ([ 26 "--python_out={}".format(out_dir.path), 27 ] + [ 28 "--proto_path={}".format(get_include_directory(i)) 29 for i in includes 30 ] + [ 31 "--proto_path={}".format(context.genfiles_dir.path), 32 ]) 33 if context.attr.plugin: 34 arguments += get_plugin_args( 35 context.executable.plugin, 36 [], 37 out_dir.path, 38 False, 39 context.attr.plugin.label.name, 40 ) 41 tools.append(context.executable.plugin) 42 43 arguments += get_proto_arguments(protos, context.genfiles_dir.path) 44 45 context.actions.run( 46 inputs = protos + includes, 47 tools = tools, 48 outputs = out_files, 49 executable = context.executable._protoc, 50 arguments = arguments, 51 mnemonic = "ProtocInvocation", 52 ) 53 54 imports = [] 55 if out_dir.import_path: 56 imports.append("%s/%s/%s" % (context.workspace_name, context.label.package, out_dir.import_path)) 57 58 return [ 59 DefaultInfo(files = depset(direct = out_files)), 60 PyInfo( 61 transitive_sources = depset(), 62 imports = depset(direct = imports), 63 ), 64 ] 65 66_generate_pb2_src = rule( 67 attrs = { 68 "deps": attr.label_list( 69 mandatory = True, 70 allow_empty = False, 71 providers = [ProtoInfo], 72 ), 73 "plugin": attr.label( 74 mandatory = False, 75 executable = True, 76 providers = ["files_to_run"], 77 cfg = "host", 78 ), 79 "_protoc": attr.label( 80 default = Label("//external:protocol_compiler"), 81 providers = ["files_to_run"], 82 executable = True, 83 cfg = "host", 84 ), 85 }, 86 implementation = _generate_py_impl, 87) 88 89def py_proto_library( 90 name, 91 deps, 92 plugin = None, 93 **kwargs): 94 """Generate python code for a protobuf. 95 96 Args: 97 name: The name of the target. 98 deps: A list of proto_library dependencies. Must contain a single element. 99 plugin: An optional custom protoc plugin to execute together with 100 generating the protobuf code. 101 **kwargs: Additional arguments to be supplied to the invocation of 102 py_library. 103 """ 104 codegen_target = "_{}_codegen".format(name) 105 if len(deps) != 1: 106 fail("Can only compile a single proto at a time.") 107 108 _generate_pb2_src( 109 name = codegen_target, 110 deps = deps, 111 plugin = plugin, 112 **kwargs 113 ) 114 115 native.py_library( 116 name = name, 117 srcs = [":{}".format(codegen_target)], 118 deps = [ 119 "@com_google_protobuf//:protobuf_python", 120 ":{}".format(codegen_target), 121 ], 122 **kwargs 123 ) 124 125def _generate_pb2_grpc_src_impl(context): 126 protos = protos_from_context(context) 127 includes = includes_from_deps(context.attr.deps) 128 out_files = declare_out_files(protos, context, _GENERATED_GRPC_PROTO_FORMAT) 129 130 plugin_flags = ["grpc_2_0"] + context.attr.strip_prefixes 131 132 arguments = [] 133 tools = [context.executable._protoc, context.executable._grpc_plugin] 134 out_dir = get_out_dir(protos, context) 135 arguments += get_plugin_args( 136 context.executable._grpc_plugin, 137 plugin_flags, 138 out_dir.path, 139 False, 140 ) 141 if context.attr.plugin: 142 arguments += get_plugin_args( 143 context.executable.plugin, 144 [], 145 out_dir.path, 146 False, 147 context.attr.plugin.label.name, 148 ) 149 tools.append(context.executable.plugin) 150 151 arguments += [ 152 "--proto_path={}".format(get_include_directory(i)) 153 for i in includes 154 ] 155 arguments += ["--proto_path={}".format(context.genfiles_dir.path)] 156 arguments += get_proto_arguments(protos, context.genfiles_dir.path) 157 158 context.actions.run( 159 inputs = protos + includes, 160 tools = tools, 161 outputs = out_files, 162 executable = context.executable._protoc, 163 arguments = arguments, 164 mnemonic = "ProtocInvocation", 165 ) 166 167 return [ 168 DefaultInfo(files = depset(direct = out_files)), 169 PyInfo( 170 transitive_sources = depset(), 171 # Imports are already configured by the generated py impl 172 imports = depset(), 173 ), 174 ] 175 176_generate_pb2_grpc_src = rule( 177 attrs = { 178 "deps": attr.label_list( 179 mandatory = True, 180 allow_empty = False, 181 providers = [ProtoInfo], 182 ), 183 "strip_prefixes": attr.string_list(), 184 "plugin": attr.label( 185 mandatory = False, 186 executable = True, 187 providers = ["files_to_run"], 188 cfg = "host", 189 ), 190 "_grpc_plugin": attr.label( 191 executable = True, 192 providers = ["files_to_run"], 193 cfg = "host", 194 default = Label("//src/compiler:grpc_python_plugin"), 195 ), 196 "_protoc": attr.label( 197 executable = True, 198 providers = ["files_to_run"], 199 cfg = "host", 200 default = Label("//external:protocol_compiler"), 201 ), 202 }, 203 implementation = _generate_pb2_grpc_src_impl, 204) 205 206def py_grpc_library( 207 name, 208 srcs, 209 deps, 210 plugin = None, 211 strip_prefixes = [], 212 **kwargs): 213 """Generate python code for gRPC services defined in a protobuf. 214 215 Args: 216 name: The name of the target. 217 srcs: (List of `labels`) a single proto_library target containing the 218 schema of the service. 219 deps: (List of `labels`) a single py_proto_library target for the 220 proto_library in `srcs`. 221 strip_prefixes: (List of `strings`) If provided, this prefix will be 222 stripped from the beginning of foo_pb2 modules imported by the 223 generated stubs. This is useful in combination with the `imports` 224 attribute of the `py_library` rule. 225 plugin: An optional custom protoc plugin to execute together with 226 generating the gRPC code. 227 **kwargs: Additional arguments to be supplied to the invocation of 228 py_library. 229 """ 230 codegen_grpc_target = "_{}_grpc_codegen".format(name) 231 if len(srcs) != 1: 232 fail("Can only compile a single proto at a time.") 233 234 if len(deps) != 1: 235 fail("Deps must have length 1.") 236 237 _generate_pb2_grpc_src( 238 name = codegen_grpc_target, 239 deps = srcs, 240 strip_prefixes = strip_prefixes, 241 plugin = plugin, 242 **kwargs 243 ) 244 245 native.py_library( 246 name = name, 247 srcs = [ 248 ":{}".format(codegen_grpc_target), 249 ], 250 deps = [ 251 Label("//src/python/grpcio/grpc:grpcio"), 252 ] + deps + [ 253 ":{}".format(codegen_grpc_target), 254 ], 255 **kwargs 256 ) 257 258def py2and3_test( 259 name, 260 py_test = native.py_test, 261 **kwargs): 262 """Runs a Python test under both Python 2 and Python 3. 263 264 Args: 265 name: The name of the test. 266 py_test: The rule to use for each test. 267 **kwargs: Keyword arguments passed directly to the underlying py_test 268 rule. 269 """ 270 if "python_version" in kwargs: 271 fail("Cannot specify 'python_version' in py2and3_test.") 272 273 names = [name + suffix for suffix in (".python2", ".python3")] 274 python_versions = ["PY2", "PY3"] 275 for case_name, python_version in zip(names, python_versions): 276 py_test( 277 name = case_name, 278 python_version = python_version, 279 **kwargs 280 ) 281 282 suite_kwargs = {} 283 if "visibility" in kwargs: 284 suite_kwargs["visibility"] = kwargs["visibility"] 285 286 native.test_suite( 287 name = name, 288 tests = names, 289 **suite_kwargs 290 ) 291