1# Copyright 2022 The Bazel Authors. All rights reserved. 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 15"""The implementation of the `py_proto_library` rule and its aspect.""" 16 17load("@rules_proto//proto:defs.bzl", "ProtoInfo", "proto_common") 18load("//python:defs.bzl", "PyInfo") 19 20PY_PROTO_TOOLCHAIN = "@rules_python//python/proto:toolchain_type" 21 22_PyProtoInfo = provider( 23 doc = "Encapsulates information needed by the Python proto rules.", 24 fields = { 25 "imports": """ 26 (depset[str]) The field forwarding PyInfo.imports coming from 27 the proto language runtime dependency.""", 28 "runfiles_from_proto_deps": """ 29 (depset[File]) Files from the transitive closure implicit proto 30 dependencies""", 31 "transitive_sources": """(depset[File]) The Python sources.""", 32 }, 33) 34 35def _filter_provider(provider, *attrs): 36 return [dep[provider] for attr in attrs for dep in attr if provider in dep] 37 38def _incompatible_toolchains_enabled(): 39 return getattr(proto_common, "INCOMPATIBLE_ENABLE_PROTO_TOOLCHAIN_RESOLUTION", False) 40 41def _py_proto_aspect_impl(target, ctx): 42 """Generates and compiles Python code for a proto_library. 43 44 The function runs protobuf compiler on the `proto_library` target generating 45 a .py file for each .proto file. 46 47 Args: 48 target: (Target) A target providing `ProtoInfo`. Usually this means a 49 `proto_library` target, but not always; you must expect to visit 50 non-`proto_library` targets, too. 51 ctx: (RuleContext) The rule context. 52 53 Returns: 54 ([_PyProtoInfo]) Providers collecting transitive information about 55 generated files. 56 """ 57 _proto_library = ctx.rule.attr 58 59 # Check Proto file names 60 for proto in target[ProtoInfo].direct_sources: 61 if proto.is_source and "-" in proto.dirname: 62 fail("Cannot generate Python code for a .proto whose path contains '-' ({}).".format( 63 proto.path, 64 )) 65 66 if _incompatible_toolchains_enabled(): 67 toolchain = ctx.toolchains[PY_PROTO_TOOLCHAIN] 68 if not toolchain: 69 fail("No toolchains registered for '%s'." % PY_PROTO_TOOLCHAIN) 70 proto_lang_toolchain_info = toolchain.proto 71 else: 72 proto_lang_toolchain_info = getattr(ctx.attr, "_aspect_proto_toolchain")[proto_common.ProtoLangToolchainInfo] 73 74 api_deps = [proto_lang_toolchain_info.runtime] 75 76 generated_sources = [] 77 proto_info = target[ProtoInfo] 78 proto_root = proto_info.proto_source_root 79 if proto_info.direct_sources: 80 # Generate py files 81 generated_sources = proto_common.declare_generated_files( 82 actions = ctx.actions, 83 proto_info = proto_info, 84 extension = "_pb2.py", 85 name_mapper = lambda name: name.replace("-", "_").replace(".", "/"), 86 ) 87 88 # Handles multiple repository and virtual import cases 89 if proto_root.startswith(ctx.bin_dir.path): 90 proto_root = proto_root[len(ctx.bin_dir.path) + 1:] 91 92 plugin_output = ctx.bin_dir.path + "/" + proto_root 93 proto_root = ctx.workspace_name + "/" + proto_root 94 95 proto_common.compile( 96 actions = ctx.actions, 97 proto_info = proto_info, 98 proto_lang_toolchain_info = proto_lang_toolchain_info, 99 generated_files = generated_sources, 100 plugin_output = plugin_output, 101 ) 102 103 # Generated sources == Python sources 104 python_sources = generated_sources 105 106 deps = _filter_provider(_PyProtoInfo, getattr(_proto_library, "deps", [])) 107 runfiles_from_proto_deps = depset( 108 transitive = [dep[DefaultInfo].default_runfiles.files for dep in api_deps] + 109 [dep.runfiles_from_proto_deps for dep in deps], 110 ) 111 transitive_sources = depset( 112 direct = python_sources, 113 transitive = [dep.transitive_sources for dep in deps], 114 ) 115 116 return [ 117 _PyProtoInfo( 118 imports = depset( 119 # Adding to PYTHONPATH so the generated modules can be 120 # imported. This is necessary when there is 121 # strip_import_prefix, the Python modules are generated under 122 # _virtual_imports. But it's undesirable otherwise, because it 123 # will put the repo root at the top of the PYTHONPATH, ahead of 124 # directories added through `imports` attributes. 125 [proto_root] if "_virtual_imports" in proto_root else [], 126 transitive = [dep[PyInfo].imports for dep in api_deps] + [dep.imports for dep in deps], 127 ), 128 runfiles_from_proto_deps = runfiles_from_proto_deps, 129 transitive_sources = transitive_sources, 130 ), 131 ] 132 133_py_proto_aspect = aspect( 134 implementation = _py_proto_aspect_impl, 135 attrs = {} if _incompatible_toolchains_enabled() else { 136 "_aspect_proto_toolchain": attr.label( 137 default = ":python_toolchain", 138 ), 139 }, 140 attr_aspects = ["deps"], 141 required_providers = [ProtoInfo], 142 provides = [_PyProtoInfo], 143 toolchains = [PY_PROTO_TOOLCHAIN] if _incompatible_toolchains_enabled() else [], 144) 145 146def _py_proto_library_rule(ctx): 147 """Merges results of `py_proto_aspect` in `deps`. 148 149 Args: 150 ctx: (RuleContext) The rule context. 151 Returns: 152 ([PyInfo, DefaultInfo, OutputGroupInfo]) 153 """ 154 if not ctx.attr.deps: 155 fail("'deps' attribute mustn't be empty.") 156 157 pyproto_infos = _filter_provider(_PyProtoInfo, ctx.attr.deps) 158 default_outputs = depset( 159 transitive = [info.transitive_sources for info in pyproto_infos], 160 ) 161 162 return [ 163 DefaultInfo( 164 files = default_outputs, 165 default_runfiles = ctx.runfiles(transitive_files = depset( 166 transitive = 167 [default_outputs] + 168 [info.runfiles_from_proto_deps for info in pyproto_infos], 169 )), 170 ), 171 OutputGroupInfo( 172 default = depset(), 173 ), 174 PyInfo( 175 transitive_sources = default_outputs, 176 imports = depset(transitive = [info.imports for info in pyproto_infos]), 177 # Proto always produces 2- and 3- compatible source files 178 has_py2_only_sources = False, 179 has_py3_only_sources = False, 180 ), 181 ] 182 183py_proto_library = rule( 184 implementation = _py_proto_library_rule, 185 doc = """ 186 Use `py_proto_library` to generate Python libraries from `.proto` files. 187 188 The convention is to name the `py_proto_library` rule `foo_py_pb2`, 189 when it is wrapping `proto_library` rule `foo_proto`. 190 191 `deps` must point to a `proto_library` rule. 192 193 Example: 194 195```starlark 196py_library( 197 name = "lib", 198 deps = [":foo_py_pb2"], 199) 200 201py_proto_library( 202 name = "foo_py_pb2", 203 deps = [":foo_proto"], 204) 205 206proto_library( 207 name = "foo_proto", 208 srcs = ["foo.proto"], 209) 210```""", 211 attrs = { 212 "deps": attr.label_list( 213 doc = """ 214 The list of `proto_library` rules to generate Python libraries for. 215 216 Usually this is just the one target: the proto library of interest. 217 It can be any target providing `ProtoInfo`.""", 218 providers = [ProtoInfo], 219 aspects = [_py_proto_aspect], 220 ), 221 }, 222 provides = [PyInfo], 223) 224