• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright (c) 2009-2021, Google LLC
2# All rights reserved.
3#
4# Use of this source code is governed by a BSD-style
5# license that can be found in the LICENSE file or at
6# https://developers.google.com/open-source/licenses/bsd
7
8"""An implementation of py_proto_library().
9
10We have to implement this ourselves because there is currently no reasonable
11py_proto_library() rule available for Bazel.
12
13Our py_proto_library() is similar to how a real py_proto_library() should work.
14But it hasn't been deeply tested or reviewed, and upb should not be in the
15business of vending py_proto_library(), so we keep it private to upb.
16"""
17
18load("@bazel_skylib//lib:paths.bzl", "paths")
19load("@rules_python//python:py_info.bzl", "PyInfo")
20load("//bazel/common:proto_info.bzl", "ProtoInfo")
21
22# Generic support code #########################################################
23
24def _get_real_short_path(file):
25    # For some reason, files from other archives have short paths that look like:
26    #   ../com_google_protobuf/google/protobuf/descriptor.proto
27    short_path = file.short_path
28    if short_path.startswith("../"):
29        second_slash = short_path.index("/", 3)
30        short_path = short_path[second_slash + 1:]
31
32    # Sometimes it has another few prefixes like:
33    #   _virtual_imports/any_proto/google/protobuf/any.proto
34    #   benchmarks/_virtual_imports/100_msgs_proto/benchmarks/100_msgs.proto
35    # We want just google/protobuf/any.proto.
36    virtual_imports = "_virtual_imports/"
37    if virtual_imports in short_path:
38        short_path = short_path.split(virtual_imports)[1].split("/", 1)[1]
39    return short_path
40
41def _get_real_root(ctx, file):
42    real_short_path = _get_real_short_path(file)
43    root = file.path[:-len(real_short_path) - 1]
44
45    if ctx.rule.attr.strip_import_prefix:
46        root = paths.join(root, ctx.rule.attr.strip_import_prefix[1:])
47
48    return root
49
50def _generate_output_file(ctx, src, extension):
51    package = ctx.label.package
52
53    strip_import_prefix = ctx.rule.attr.strip_import_prefix
54    if strip_import_prefix and strip_import_prefix != "/":
55        if not package.startswith(strip_import_prefix[1:]):
56            fail("%s does not begin with prefix %s" % (package, strip_import_prefix))
57        package = package[len(strip_import_prefix):]
58
59    real_short_path = _get_real_short_path(src)
60    real_short_path = paths.relativize(real_short_path, package)
61    output_filename = paths.replace_extension(real_short_path, extension)
62    ret = ctx.actions.declare_file(output_filename)
63    return ret
64
65# py_proto_library() ###########################################################
66
67def _py_proto_library_rule_impl(ctx):
68    # A real py_proto_library() should enforce this constraint.
69    # We don't bother for now, since it saves us some effort not to.
70    #
71    # if len(ctx.attr.deps) != 1:
72    #     fail("only one deps dependency allowed.")
73
74    files = []
75    for dep in ctx.attr.deps:
76        files += dep[PyInfo].transitive_sources.to_list()
77    return [
78        DefaultInfo(files = depset(direct = files)),
79    ]
80
81def _py_proto_library_aspect_impl(target, ctx):
82    proto_info = target[ProtoInfo]
83    proto_sources = proto_info.direct_sources
84    srcs = [_generate_output_file(ctx, name, "_pb2.py") for name in proto_sources]
85    transitive_sets = proto_info.transitive_descriptor_sets.to_list()
86    ctx.actions.run(
87        inputs = depset(
88            direct = [proto_info.direct_descriptor_set],
89            transitive = [proto_info.transitive_descriptor_sets],
90        ),
91        outputs = srcs,
92        executable = ctx.executable._protoc,
93        arguments = [
94                        "--python_out=" + _get_real_root(ctx, srcs[0]),
95                        "--descriptor_set_in=" + ctx.configuration.host_path_separator.join([f.path for f in transitive_sets]),
96                    ] +
97                    [_get_real_short_path(file) for file in proto_sources],
98        progress_message = "Generating Python protos for :" + ctx.label.name,
99    )
100    outs_depset = depset(srcs)
101    return [
102        PyInfo(transitive_sources = outs_depset),
103    ]
104
105_py_proto_library_aspect = aspect(
106    attrs = {
107        "_protoc": attr.label(
108            executable = True,
109            cfg = "exec",
110            default = "//:protoc",
111        ),
112    },
113    implementation = _py_proto_library_aspect_impl,
114    provides = [
115        PyInfo,
116    ],
117    attr_aspects = ["deps"],
118)
119
120py_proto_library = rule(
121    implementation = _py_proto_library_rule_impl,
122    attrs = {
123        "deps": attr.label_list(
124            aspects = [_py_proto_library_aspect],
125            allow_rules = ["proto_library"],
126            providers = [ProtoInfo],
127        ),
128    },
129)
130