blob: 86af999502375fcf61486048d13e27b3a6cf841b [file] [log] [blame]
"""The implementation of the `py_proto_library` rule and its aspect."""
load("@rules_python//python:py_info.bzl", "PyInfo")
load("//bazel/common:proto_common.bzl", "proto_common")
load("//bazel/common:proto_info.bzl", "ProtoInfo")
PY_PROTO_TOOLCHAIN = "@rules_python//python/proto:toolchain_type"
_PyProtoInfo = provider(
doc = "Encapsulates information needed by the Python proto rules.",
fields = {
"imports": """
(depset[str]) The field forwarding PyInfo.imports coming from
the proto language runtime dependency.""",
"runfiles_from_proto_deps": """
(depset[File]) Files from the transitive closure implicit proto
dependencies""",
"transitive_sources": """(depset[File]) The Python sources.""",
},
)
def _filter_provider(provider, *attrs):
return [dep[provider] for attr in attrs for dep in attr if provider in dep]
def _incompatible_toolchains_enabled():
return getattr(proto_common, "INCOMPATIBLE_ENABLE_PROTO_TOOLCHAIN_RESOLUTION", False)
def _py_proto_aspect_impl(target, ctx):
"""Generates and compiles Python code for a proto_library.
The function runs protobuf compiler on the `proto_library` target generating
a .py file for each .proto file.
Args:
target: (Target) A target providing `ProtoInfo`. Usually this means a
`proto_library` target, but not always; you must expect to visit
non-`proto_library` targets, too.
ctx: (RuleContext) The rule context.
Returns:
([_PyProtoInfo]) Providers collecting transitive information about
generated files.
"""
_proto_library = ctx.rule.attr
# Check Proto file names
for proto in target[ProtoInfo].direct_sources:
if proto.is_source and "-" in proto.dirname:
fail("Cannot generate Python code for a .proto whose path contains '-' ({}).".format(
proto.path,
))
if _incompatible_toolchains_enabled():
toolchain = ctx.toolchains[PY_PROTO_TOOLCHAIN]
if not toolchain:
fail("No toolchains registered for '%s'." % PY_PROTO_TOOLCHAIN)
proto_lang_toolchain_info = toolchain.proto
else:
proto_lang_toolchain_info = getattr(ctx.attr, "_aspect_proto_toolchain")[proto_common.ProtoLangToolchainInfo]
api_deps = [proto_lang_toolchain_info.runtime]
generated_sources = []
proto_info = target[ProtoInfo]
proto_root = proto_info.proto_source_root
if proto_info.direct_sources:
# Generate py files
generated_sources = proto_common.declare_generated_files(
actions = ctx.actions,
proto_info = proto_info,
extension = "_pb2.py",
name_mapper = lambda name: name.replace("-", "_").replace(".", "/"),
)
# Handles multiple repository and virtual import cases
if proto_root.startswith(ctx.bin_dir.path):
proto_root = proto_root[len(ctx.bin_dir.path) + 1:]
plugin_output = ctx.bin_dir.path + "/" + proto_root
proto_root = ctx.workspace_name + "/" + proto_root
proto_common.compile(
actions = ctx.actions,
proto_info = proto_info,
proto_lang_toolchain_info = proto_lang_toolchain_info,
generated_files = generated_sources,
plugin_output = plugin_output,
)
# Generated sources == Python sources
python_sources = generated_sources
deps = _filter_provider(_PyProtoInfo, getattr(_proto_library, "deps", []))
runfiles_from_proto_deps = depset(
transitive = [dep[DefaultInfo].default_runfiles.files for dep in api_deps] +
[dep.runfiles_from_proto_deps for dep in deps],
)
transitive_sources = depset(
direct = python_sources,
transitive = [dep.transitive_sources for dep in deps],
)
return [
_PyProtoInfo(
imports = depset(
# Adding to PYTHONPATH so the generated modules can be
# imported. This is necessary when there is
# strip_import_prefix, the Python modules are generated under
# _virtual_imports. But it's undesirable otherwise, because it
# will put the repo root at the top of the PYTHONPATH, ahead of
# directories added through `imports` attributes.
[proto_root] if "_virtual_imports" in proto_root else [],
transitive = [dep[PyInfo].imports for dep in api_deps] + [dep.imports for dep in deps],
),
runfiles_from_proto_deps = runfiles_from_proto_deps,
transitive_sources = transitive_sources,
),
]
_py_proto_aspect = aspect(
implementation = _py_proto_aspect_impl,
attrs = {} if _incompatible_toolchains_enabled() else {
"_aspect_proto_toolchain": attr.label(
default = "//python:python_toolchain",
),
},
attr_aspects = ["deps"],
required_providers = [ProtoInfo],
provides = [_PyProtoInfo],
toolchains = [PY_PROTO_TOOLCHAIN] if _incompatible_toolchains_enabled() else [],
)
def _py_proto_library_rule(ctx):
"""Merges results of `py_proto_aspect` in `deps`.
Args:
ctx: (RuleContext) The rule context.
Returns:
([PyInfo, DefaultInfo, OutputGroupInfo])
"""
if not ctx.attr.deps:
fail("'deps' attribute mustn't be empty.")
pyproto_infos = _filter_provider(_PyProtoInfo, ctx.attr.deps)
default_outputs = depset(
transitive = [info.transitive_sources for info in pyproto_infos],
)
return [
DefaultInfo(
files = default_outputs,
default_runfiles = ctx.runfiles(transitive_files = depset(
transitive =
[default_outputs] +
[info.runfiles_from_proto_deps for info in pyproto_infos],
)),
),
OutputGroupInfo(
default = depset(),
),
PyInfo(
transitive_sources = default_outputs,
imports = depset(transitive = [info.imports for info in pyproto_infos]),
# Proto always produces 2- and 3- compatible source files
has_py2_only_sources = False,
has_py3_only_sources = False,
),
]
py_proto_library = rule(
implementation = _py_proto_library_rule,
doc = """
Use `py_proto_library` to generate Python libraries from `.proto` files.
The convention is to name the `py_proto_library` rule `foo_py_pb2`,
when it is wrapping `proto_library` rule `foo_proto`.
`deps` must point to a `proto_library` rule.
Example:
```starlark
py_library(
name = "lib",
deps = [":foo_py_pb2"],
)
py_proto_library(
name = "foo_py_pb2",
deps = [":foo_proto"],
)
proto_library(
name = "foo_proto",
srcs = ["foo.proto"],
)
```""",
attrs = {
"deps": attr.label_list(
doc = """
The list of `proto_library` rules to generate Python libraries for.
Usually this is just the one target: the proto library of interest.
It can be any target providing `ProtoInfo`.""",
providers = [ProtoInfo],
aspects = [_py_proto_aspect],
),
},
provides = [PyInfo],
)