xref: /aosp_15_r20/external/bazelbuild-rules_python/python/private/proto/py_proto_library.bzl (revision 60517a1edbc8ecf509223e9af94a7adec7d736b8)
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