1# Copyright 2021 The 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"""Generates and compiles Python gRPC stubs from proto_library rules.""" 15 16load("@rules_proto//proto:defs.bzl", "ProtoInfo") 17load( 18 "//bazel:protobuf.bzl", 19 "declare_out_files", 20 "get_include_directory", 21 "get_out_dir", 22 "get_plugin_args", 23 "get_proto_arguments", 24 "get_staged_proto_file", 25 "includes_from_deps", 26 "is_well_known", 27 "protos_from_context", 28) 29 30_GENERATED_PROTO_FORMAT = "{}_pb2.py" 31_GENERATED_PROTO_STUB_FORMAT = "{}_pb2.pyi" 32_GENERATED_GRPC_PROTO_FORMAT = "{}_pb2_grpc.py" 33 34PyProtoInfo = provider( 35 "The Python outputs from the Protobuf compiler.", 36 fields = { 37 "py_info": "A PyInfo provider for the generated code.", 38 "generated_py_srcs": "The direct (not transitive) generated Python source files.", 39 }, 40) 41 42def _merge_pyinfos(pyinfos): 43 return PyInfo( 44 transitive_sources = depset(transitive = [p.transitive_sources for p in pyinfos]), 45 imports = depset(transitive = [p.imports for p in pyinfos]), 46 ) 47 48def _gen_py_aspect_impl(target, context): 49 # Early return for well-known protos. 50 if is_well_known(str(context.label)): 51 return [ 52 PyProtoInfo( 53 py_info = context.attr._protobuf_library[PyInfo], 54 generated_py_srcs = [], 55 ), 56 ] 57 58 protos = [] 59 for p in target[ProtoInfo].direct_sources: 60 protos.append(get_staged_proto_file(target.label, context, p)) 61 62 includes = depset(direct = protos, transitive = [target[ProtoInfo].transitive_imports]) 63 out_files = (declare_out_files(protos, context, _GENERATED_PROTO_FORMAT) + 64 declare_out_files(protos, context, _GENERATED_PROTO_STUB_FORMAT)) 65 generated_py_srcs = out_files 66 67 tools = [context.executable._protoc] 68 69 out_dir = get_out_dir(protos, context) 70 71 arguments = ([ 72 "--python_out={}".format(out_dir.path), 73 "--pyi_out={}".format(out_dir.path), 74 ] + [ 75 "--proto_path={}".format(get_include_directory(i)) 76 for i in includes.to_list() 77 ] + [ 78 "--proto_path={}".format(context.genfiles_dir.path), 79 ]) 80 81 arguments += get_proto_arguments(protos, context.genfiles_dir.path) 82 83 context.actions.run( 84 inputs = protos + includes.to_list(), 85 tools = tools, 86 outputs = out_files, 87 executable = context.executable._protoc, 88 arguments = arguments, 89 mnemonic = "ProtocInvocation", 90 ) 91 92 imports = [] 93 if out_dir.import_path: 94 imports.append("{}/{}".format(context.workspace_name, out_dir.import_path)) 95 96 py_info = PyInfo(transitive_sources = depset(direct = out_files), imports = depset(direct = imports)) 97 return PyProtoInfo( 98 py_info = _merge_pyinfos( 99 [ 100 py_info, 101 context.attr._protobuf_library[PyInfo], 102 ] + [dep[PyProtoInfo].py_info for dep in context.rule.attr.deps], 103 ), 104 generated_py_srcs = generated_py_srcs, 105 ) 106 107_gen_py_aspect = aspect( 108 implementation = _gen_py_aspect_impl, 109 attr_aspects = ["deps"], 110 fragments = ["py"], 111 attrs = { 112 "_protoc": attr.label( 113 default = Label("//external:protocol_compiler"), 114 providers = ["files_to_run"], 115 executable = True, 116 cfg = "exec", 117 ), 118 "_protobuf_library": attr.label( 119 default = Label("@com_google_protobuf//:protobuf_python"), 120 providers = [PyInfo], 121 ), 122 }, 123) 124 125def _generate_py_impl(context): 126 if (len(context.attr.deps) != 1): 127 fail("Can only compile a single proto at a time.") 128 129 py_sources = [] 130 131 # If the proto_library this rule *directly* depends on is in another 132 # package, then we generate .py files to import them in this package. This 133 # behavior is needed to allow rearranging of import paths to make Bazel 134 # outputs align with native python workflows. 135 # 136 # Note that this approach is vulnerable to protoc defining __all__ or other 137 # symbols with __ prefixes that need to be directly imported. Since these 138 # names are likely to be reserved for private APIs, the risk is minimal. 139 if context.label.package != context.attr.deps[0].label.package: 140 for py_src in context.attr.deps[0][PyProtoInfo].generated_py_srcs: 141 reimport_py_file = context.actions.declare_file(py_src.basename) 142 py_sources.append(reimport_py_file) 143 import_line = "from %s import *" % py_src.short_path.replace("..", "external").replace("/", ".")[:-len(".py")] 144 context.actions.write(reimport_py_file, import_line) 145 146 # Collect output PyInfo provider. 147 imports = [context.label.package + "/" + i for i in context.attr.imports] 148 py_info = PyInfo(transitive_sources = depset(direct = py_sources), imports = depset(direct = imports)) 149 out_pyinfo = _merge_pyinfos([py_info, context.attr.deps[0][PyProtoInfo].py_info]) 150 151 runfiles = context.runfiles(files = out_pyinfo.transitive_sources.to_list()).merge(context.attr._protobuf_library[DefaultInfo].data_runfiles) 152 return [ 153 DefaultInfo( 154 files = out_pyinfo.transitive_sources, 155 runfiles = runfiles, 156 ), 157 out_pyinfo, 158 ] 159 160py_proto_library = rule( 161 attrs = { 162 "deps": attr.label_list( 163 mandatory = True, 164 allow_empty = False, 165 providers = [ProtoInfo], 166 aspects = [_gen_py_aspect], 167 ), 168 "_protoc": attr.label( 169 default = Label("//external:protocol_compiler"), 170 providers = ["files_to_run"], 171 executable = True, 172 cfg = "exec", 173 ), 174 "_protobuf_library": attr.label( 175 default = Label("@com_google_protobuf//:protobuf_python"), 176 providers = [PyInfo], 177 ), 178 "imports": attr.string_list(), 179 }, 180 implementation = _generate_py_impl, 181) 182 183def _generate_pb2_grpc_src_impl(context): 184 protos = protos_from_context(context) 185 includes = includes_from_deps(context.attr.deps) 186 out_files = declare_out_files(protos, context, _GENERATED_GRPC_PROTO_FORMAT) 187 188 plugin_flags = ["grpc_2_0"] + context.attr.strip_prefixes 189 190 arguments = [] 191 tools = [context.executable._protoc, context.executable._grpc_plugin] 192 out_dir = get_out_dir(protos, context) 193 if out_dir.import_path: 194 # is virtual imports 195 out_path = out_dir.path 196 else: 197 out_path = context.genfiles_dir.path 198 arguments += get_plugin_args( 199 context.executable._grpc_plugin, 200 plugin_flags, 201 out_path, 202 False, 203 ) 204 205 arguments += [ 206 "--proto_path={}".format(get_include_directory(i)) 207 for i in includes 208 ] 209 arguments.append("--proto_path={}".format(context.genfiles_dir.path)) 210 arguments += get_proto_arguments(protos, context.genfiles_dir.path) 211 212 context.actions.run( 213 inputs = protos + includes, 214 tools = tools, 215 outputs = out_files, 216 executable = context.executable._protoc, 217 arguments = arguments, 218 mnemonic = "ProtocInvocation", 219 ) 220 221 p = PyInfo(transitive_sources = depset(direct = out_files)) 222 py_info = _merge_pyinfos( 223 [ 224 p, 225 context.attr.grpc_library[PyInfo], 226 ] + [dep[PyInfo] for dep in context.attr.py_deps], 227 ) 228 229 runfiles = context.runfiles(files = out_files, transitive_files = py_info.transitive_sources).merge(context.attr.grpc_library[DefaultInfo].data_runfiles) 230 231 return [ 232 DefaultInfo( 233 files = depset(direct = out_files), 234 runfiles = runfiles, 235 ), 236 py_info, 237 ] 238 239_generate_pb2_grpc_src = rule( 240 attrs = { 241 "deps": attr.label_list( 242 mandatory = True, 243 allow_empty = False, 244 providers = [ProtoInfo], 245 ), 246 "py_deps": attr.label_list( 247 mandatory = True, 248 allow_empty = False, 249 providers = [PyInfo], 250 ), 251 "strip_prefixes": attr.string_list(), 252 "_grpc_plugin": attr.label( 253 executable = True, 254 providers = ["files_to_run"], 255 cfg = "exec", 256 default = Label("//src/compiler:grpc_python_plugin"), 257 ), 258 "_protoc": attr.label( 259 executable = True, 260 providers = ["files_to_run"], 261 cfg = "exec", 262 default = Label("//external:protocol_compiler"), 263 ), 264 "grpc_library": attr.label( 265 default = Label("//src/python/grpcio/grpc:grpcio"), 266 providers = [PyInfo], 267 ), 268 }, 269 implementation = _generate_pb2_grpc_src_impl, 270) 271 272def py_grpc_library( 273 name, 274 srcs, 275 deps, 276 strip_prefixes = [], 277 grpc_library = Label("//src/python/grpcio/grpc:grpcio"), 278 **kwargs): 279 """Generate python code for gRPC services defined in a protobuf. 280 281 Args: 282 name: The name of the target. 283 srcs: (List of `labels`) a single proto_library target containing the 284 schema of the service. 285 deps: (List of `labels`) a single py_proto_library target for the 286 proto_library in `srcs`. 287 strip_prefixes: (List of `strings`) If provided, this prefix will be 288 stripped from the beginning of foo_pb2 modules imported by the 289 generated stubs. This is useful in combination with the `imports` 290 attribute of the `py_library` rule. 291 grpc_library: (`label`) a single `py_library` target representing the 292 python gRPC library target to be depended upon. This can be used to 293 generate code that depends on `grpcio` from the Python Package Index. 294 **kwargs: Additional arguments to be supplied to the invocation of 295 py_library. 296 """ 297 if len(srcs) != 1: 298 fail("Can only compile a single proto at a time.") 299 300 if len(deps) != 1: 301 fail("Deps must have length 1.") 302 303 _generate_pb2_grpc_src( 304 name = name, 305 deps = srcs, 306 py_deps = deps, 307 strip_prefixes = strip_prefixes, 308 grpc_library = grpc_library, 309 **kwargs 310 ) 311