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("__main__/%s" % 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 imports = [] 168 if out_dir.import_path: 169 imports.append("__main__/%s" % out_dir.import_path) 170 171 return [ 172 DefaultInfo(files = depset(direct = out_files)), 173 PyInfo( 174 transitive_sources = depset(), 175 imports = depset(direct = imports), 176 ), 177 ] 178 179_generate_pb2_grpc_src = rule( 180 attrs = { 181 "deps": attr.label_list( 182 mandatory = True, 183 allow_empty = False, 184 providers = [ProtoInfo], 185 ), 186 "strip_prefixes": attr.string_list(), 187 "plugin": attr.label( 188 mandatory = False, 189 executable = True, 190 providers = ["files_to_run"], 191 cfg = "host", 192 ), 193 "_grpc_plugin": attr.label( 194 executable = True, 195 providers = ["files_to_run"], 196 cfg = "host", 197 default = Label("//src/compiler:grpc_python_plugin"), 198 ), 199 "_protoc": attr.label( 200 executable = True, 201 providers = ["files_to_run"], 202 cfg = "host", 203 default = Label("//external:protocol_compiler"), 204 ), 205 }, 206 implementation = _generate_pb2_grpc_src_impl, 207) 208 209def py_grpc_library( 210 name, 211 srcs, 212 deps, 213 plugin = None, 214 strip_prefixes = [], 215 **kwargs): 216 """Generate python code for gRPC services defined in a protobuf. 217 218 Args: 219 name: The name of the target. 220 srcs: (List of `labels`) a single proto_library target containing the 221 schema of the service. 222 deps: (List of `labels`) a single py_proto_library target for the 223 proto_library in `srcs`. 224 strip_prefixes: (List of `strings`) If provided, this prefix will be 225 stripped from the beginning of foo_pb2 modules imported by the 226 generated stubs. This is useful in combination with the `imports` 227 attribute of the `py_library` rule. 228 plugin: An optional custom protoc plugin to execute together with 229 generating the gRPC code. 230 **kwargs: Additional arguments to be supplied to the invocation of 231 py_library. 232 """ 233 codegen_grpc_target = "_{}_grpc_codegen".format(name) 234 if len(srcs) != 1: 235 fail("Can only compile a single proto at a time.") 236 237 if len(deps) != 1: 238 fail("Deps must have length 1.") 239 240 _generate_pb2_grpc_src( 241 name = codegen_grpc_target, 242 deps = srcs, 243 strip_prefixes = strip_prefixes, 244 plugin = plugin, 245 **kwargs 246 ) 247 248 native.py_library( 249 name = name, 250 srcs = [ 251 ":{}".format(codegen_grpc_target), 252 ], 253 deps = [ 254 Label("//src/python/grpcio/grpc:grpcio"), 255 ] + deps + [ 256 ":{}".format(codegen_grpc_target), 257 ], 258 **kwargs 259 ) 260 261def py2and3_test( 262 name, 263 py_test = native.py_test, 264 **kwargs): 265 """Runs a Python test under both Python 2 and Python 3. 266 267 Args: 268 name: The name of the test. 269 py_test: The rule to use for each test. 270 **kwargs: Keyword arguments passed directly to the underlying py_test 271 rule. 272 """ 273 if "python_version" in kwargs: 274 fail("Cannot specify 'python_version' in py2and3_test.") 275 276 names = [name + suffix for suffix in (".python2", ".python3")] 277 python_versions = ["PY2", "PY3"] 278 for case_name, python_version in zip(names, python_versions): 279 py_test( 280 name = case_name, 281 python_version = python_version, 282 **kwargs 283 ) 284 285 suite_kwargs = {} 286 if "visibility" in kwargs: 287 suite_kwargs["visibility"] = kwargs["visibility"] 288 289 native.test_suite( 290 name = name, 291 tests = names, 292 **suite_kwargs 293 ) 294