|
| 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 | + |
| 17 | +load("@rules_proto//proto:defs.bzl", "ProtoInfo", "proto_common") |
| 18 | +load("//python:defs.bzl", "PyInfo") |
| 19 | + |
| 20 | +ProtoLangToolchainInfo = proto_common.ProtoLangToolchainInfo |
| 21 | + |
| 22 | +_PyProtoInfo = provider( |
| 23 | + doc = "Encapsulates information needed by the Python proto rules.", |
| 24 | + fields = { |
| 25 | + "runfiles_from_proto_deps": """ |
| 26 | + (depset[File]) Files from the transitive closure implicit proto |
| 27 | + dependencies""", |
| 28 | + "transitive_sources": """(depset[File]) The Python sources.""", |
| 29 | + }, |
| 30 | +) |
| 31 | + |
| 32 | +def _filter_provider(provider, *attrs): |
| 33 | + return [dep[provider] for attr in attrs for dep in attr if provider in dep] |
| 34 | + |
| 35 | +def _py_proto_aspect_impl(target, ctx): |
| 36 | + """Generates and compiles Python code for a proto_library. |
| 37 | +
|
| 38 | + The function runs protobuf compiler on the `proto_library` target generating |
| 39 | + a .py file for each .proto file. |
| 40 | +
|
| 41 | + Args: |
| 42 | + target: (Target) A target providing `ProtoInfo`. Usually this means a |
| 43 | + `proto_library` target, but not always; you must expect to visit |
| 44 | + non-`proto_library` targets, too. |
| 45 | + ctx: (RuleContext) The rule context. |
| 46 | +
|
| 47 | + Returns: |
| 48 | + ([_PyProtoInfo]) Providers collecting transitive information about |
| 49 | + generated files. |
| 50 | + """ |
| 51 | + |
| 52 | + _proto_library = ctx.rule.attr |
| 53 | + |
| 54 | + # Check Proto file names |
| 55 | + for proto in target[ProtoInfo].direct_sources: |
| 56 | + if proto.is_source and "-" in proto.dirname: |
| 57 | + fail("Cannot generate Python code for a .proto whose path contains '-' ({}).".format( |
| 58 | + proto.path, |
| 59 | + )) |
| 60 | + |
| 61 | + proto_lang_toolchain_info = ctx.attr._aspect_proto_toolchain[ProtoLangToolchainInfo] |
| 62 | + api_deps = [proto_lang_toolchain_info.runtime] |
| 63 | + |
| 64 | + generated_sources = [] |
| 65 | + proto_info = target[ProtoInfo] |
| 66 | + if proto_info.direct_sources: |
| 67 | + # Generate py files |
| 68 | + generated_sources = proto_common.declare_generated_files( |
| 69 | + actions = ctx.actions, |
| 70 | + proto_info = proto_info, |
| 71 | + extension = "_pb2.py", |
| 72 | + name_mapper = lambda name: name.replace("-", "_").replace(".", "/"), |
| 73 | + ) |
| 74 | + |
| 75 | + proto_common.compile( |
| 76 | + actions = ctx.actions, |
| 77 | + proto_info = proto_info, |
| 78 | + proto_lang_toolchain_info = proto_lang_toolchain_info, |
| 79 | + generated_files = generated_sources, |
| 80 | + plugin_output = ctx.bin_dir.path, |
| 81 | + ) |
| 82 | + |
| 83 | + # Generated sources == Python sources |
| 84 | + python_sources = generated_sources |
| 85 | + |
| 86 | + deps = _filter_provider(_PyProtoInfo, getattr(_proto_library, "deps", [])) |
| 87 | + runfiles_from_proto_deps = depset( |
| 88 | + transitive = [dep[DefaultInfo].default_runfiles.files for dep in
1241
api_deps] + |
| 89 | + [dep.runfiles_from_proto_deps for dep in deps], |
| 90 | + ) |
| 91 | + transitive_sources = depset( |
| 92 | + direct = python_sources, |
| 93 | + transitive = [dep.transitive_sources for dep in deps], |
| 94 | + ) |
| 95 | + |
| 96 | + return [ |
| 97 | + _PyProtoInfo( |
| 98 | + runfiles_from_proto_deps = runfiles_from_proto_deps, |
| 99 | + transitive_sources = transitive_sources, |
| 100 | + ), |
| 101 | + ] |
| 102 | + |
| 103 | +_py_proto_aspect = aspect( |
| 104 | + implementation = _py_proto_aspect_impl, |
| 105 | + attrs = { |
| 106 | + "_aspect_proto_toolchain": attr.label( |
| 107 | + default = ":python_toolchain", |
| 108 | + ), |
| 109 | + }, |
| 110 | + attr_aspects = ["deps"], |
| 111 | + required_providers = [ProtoInfo], |
| 112 | + provides = [_PyProtoInfo], |
| 113 | +) |
| 114 | + |
| 115 | +def _py_proto_library_rule(ctx): |
| 116 | + """Merges results of `py_proto_aspect` in `deps`. |
| 117 | +
|
| 118 | + Args: |
| 119 | + ctx: (RuleContext) The rule context. |
| 120 | + Returns: |
| 121 | + ([PyInfo, DefaultInfo, OutputGroupInfo]) |
| 122 | + """ |
| 123 | + if not ctx.attr.deps: |
| 124 | + fail("'deps' attribute mustn't be empty.") |
| 125 | + |
| 126 | + pyproto_infos = _filter_provider(_PyProtoInfo, ctx.attr.deps) |
| 127 | + default_outputs = depset( |
| 128 | + transitive = [info.transitive_sources for info in pyproto_infos], |
| 129 | + ) |
| 130 | + |
| 131 | + return [ |
| 132 | + DefaultInfo( |
| 133 | + files = default_outputs, |
| 134 | + default_runfiles = ctx.runfiles(transitive_files = depset( |
| 135 | + transitive = |
| 136 | + [default_outputs] + |
| 137 | + [info.runfiles_from_proto_deps for info in pyproto_infos], |
| 138 | + )), |
| 139 | + ), |
| 140 | + OutputGroupInfo( |
| 141 | + default = depset(), |
| 142 | + ), |
| 143 | + PyInfo( |
| 144 | + transitive_sources = default_outputs, |
| 145 | + # Proto always produces 2- and 3- compatible source files |
| 146 | + has_py2_only_sources = False, |
| 147 | + has_py3_only_sources = False, |
| 148 | + ), |
| 149 | + ] |
| 150 | + |
| 151 | +py_proto_library = rule( |
| 152 | + implementation = _py_proto_library_rule, |
| 153 | + doc = """ |
| 154 | + Use `py_proto_library` to generate Python libraries from `.proto` files. |
| 155 | +
|
| 156 | + The convention is to name the `py_proto_library` rule `foo_py_pb2`, |
| 157 | + when it is wrapping `proto_library` rule `foo_proto`. |
| 158 | +
|
| 159 | + `deps` must point to a `proto_library` rule. |
| 160 | +
|
| 161 | + Example: |
| 162 | +
|
| 163 | +```starlark |
| 164 | +py_library( |
| 165 | + name = "lib", |
| 166 | + deps = [":foo_py_pb2"], |
| 167 | +) |
| 168 | +
|
| 169 | +py_proto_library( |
| 170 | + name = "foo_py_pb2", |
| 171 | + deps = [":foo_proto"], |
| 172 | +) |
| 173 | +
|
| 174 | +proto_library( |
| 175 | + name = "foo_proto", |
| 176 | + srcs = ["foo.proto"], |
| 177 | +) |
| 178 | +```""", |
| 179 | + attrs = { |
| 180 | + "deps": attr.label_list( |
| 181 | + doc = """ |
| 182 | + The list of `proto_library` rules to generate Python libraries for. |
| 183 | +
|
| 184 | + Usually this is just the one target: the proto library of interest. |
| 185 | + It can be any target providing `ProtoInfo`.""", |
| 186 | + providers = [ProtoInfo], |
| 187 | + aspects = [_py_proto_aspect], |
| 188 | + ), |
| 189 | + }, |
| 190 | + provides = [PyInfo], |
| 191 | +) |
0 commit comments