"""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], )