rules/python/py_proto.bzl

"""
Copyright (C) 2022 The Android Open Source Project

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

load("@bazel_skylib//lib:paths.bzl", "paths")

def _py_proto_sources_gen_rule_impl(ctx):
    imports = []
    all_outputs = []
    for dep in ctx.attr.deps:
        proto_info = dep[ProtoInfo]

        outputs = []
        for name in proto_info.direct_sources:
            outputs.append(ctx.actions.declare_file(paths.replace_extension(name.basename, "_pb2.py"), sibling = name))

        args = ctx.actions.args()
        args.add("--python_out=" + proto_info.proto_source_root)
        args.add_all(["-I", proto_info.proto_source_root])
        args.add_all(proto_info.direct_sources)

        if proto_info.proto_source_root != ".":
            imports.append(paths.join("__main__", paths.relativize(proto_info.proto_source_root, ctx.bin_dir.path)))

        # It's not clear what to do with transititve imports/sources
        if len(proto_info.transitive_imports.to_list()) > len(proto_info.direct_sources) or len(proto_info.transitive_sources.to_list()) > len(proto_info.direct_sources):
            fail("TODO: Transitive imports/sources of python protos")

        ctx.actions.run(
            inputs = depset(
                direct = proto_info.direct_sources,
                transitive = [proto_info.transitive_imports],
            ),
            executable = ctx.executable._protoc,
            outputs = outputs,
            arguments = [args],
            mnemonic = "PyProtoGen",
        )

        all_outputs.extend(outputs)

    output_depset = depset(direct = all_outputs)
    return [
        DefaultInfo(files = output_depset),
        PyInfo(
            transitive_sources = output_depset,
            # If proto_source_root is set to something other than the root of the workspace, import the current package.
            # It's always the current package because it's the path to where we generated the python sources, not to where
            # the proto sources are.
            imports = depset(direct = imports),
        ),
    ]

_py_proto_sources_gen = rule(
    implementation = _py_proto_sources_gen_rule_impl,
    attrs = {
        "deps": attr.label_list(
            providers = [ProtoInfo],
            doc = "proto_library or any other target exposing ProtoInfo provider with *.proto files",
            mandatory = True,
        ),
        "_protoc": attr.label(
            default = Label("//external/protobuf:aprotoc"),
            executable = True,
            cfg = "exec",
        ),
    },
)

def py_proto_library(
        name,
        deps = [],
        target_compatible_with = [],
        **kwargs):
    proto_lib_name = name + "_proto_gen"

    _py_proto_sources_gen(
        name = proto_lib_name,
        deps = deps,
        **kwargs
    )

    # There may be a better way to do this, but proto_lib_name appears in both srcs
    # and deps because it must appear in srcs to cause the protobuf files to
    # actually be compiled, and it must appear in deps for the PyInfo provider to
    # be respected and the "imports" path to be included in this library.
    native.py_library(
        name = name,
        srcs = [":" + proto_lib_name],
        deps = [":" + proto_lib_name] + (["//external/protobuf:libprotobuf-python"] if "libprotobuf-python" not in name else []),
        target_compatible_with = target_compatible_with,
        **kwargs
    )