|
| 1 | +"""The implementation of the `py_proto_library` rule and its aspect.""" |
| 2 | + |
| 3 | +load("@rules_python//python:py_info.bzl", "PyInfo") |
| 4 | +load("//bazel/common:proto_common.bzl", "proto_common") |
| 5 | +load("//bazel/common:proto_info.bzl", "ProtoInfo") |
| 6 | + |
| 7 | +ProtoLangToolchainInfo = proto_common.ProtoLangToolchainInfo |
| 8 | + |
| 9 | +_PyProtoInfo = provider( |
| 10 | + doc = "Encapsulates information needed by the Python proto rules.", |
| 11 | + fields = { |
| 12 | + "imports": """ |
| 13 | + (depset[str]) The field forwarding PyInfo.imports coming from |
| 14 | + the proto language runtime dependency.""", |
| 15 | + "runfiles_from_proto_deps": """ |
| 16 | + (depset[File]) Files from the transitive closure implicit proto |
| 17 | + dependencies""", |
| 18 | + "transitive_sources": """(depset[File]) The Python sources.""", |
| 19 | + }, |
| 20 | +) |
| 21 | + |
| 22 | +def _filter_provider(provider, *attrs): |
| 23 | + return [dep[provider] for attr in attrs for dep in attr if provider in dep] |
| 24 | + |
| 25 | +def _py_proto_aspect_impl(target, ctx): |
| 26 | + """Generates and compiles Python code for a proto_library. |
| 27 | +
|
| 28 | + The function runs protobuf compiler on the `proto_library` target generating |
| 29 | + a .py file for each .proto file. |
| 30 | +
|
| 31 | + Args: |
| 32 | + target: (Target) A target providing `ProtoInfo`. Usually this means a |
| 33 | + `proto_library` target, but not always; you must expect to visit |
| 34 | + non-`proto_library` targets, too. |
| 35 | + ctx: (RuleContext) The rule context. |
| 36 | +
|
| 37 | + Returns: |
| 38 | + ([_PyProtoInfo]) Providers collecting transitive information about |
| 39 | + generated files. |
| 40 | + """ |
| 41 | + |
| 42 | + _proto_library = ctx.rule.attr |
| 43 | + |
| 44 | + # Check Proto file names |
| 45 | + for proto in target[ProtoInfo].direct_sources: |
| 46 | + if proto.is_source and "-" in proto.dirname: |
| 47 | + fail("Cannot generate Python code for a .proto whose path contains '-' ({}).".format( |
| 48 | + proto.path, |
| 49 | + )) |
| 50 | + |
| 51 | + proto_lang_toolchain_info = ctx.attr._aspect_proto_toolchain[ProtoLangToolchainInfo] |
| 52 | + api_deps = [proto_lang_toolchain_info.runtime] |
| 53 | + |
| 54 | + generated_sources = [] |
| 55 | + proto_info = target[ProtoInfo] |
| 56 | + proto_root = proto_info.proto_source_root |
| 57 | + if proto_info.direct_sources: |
| 58 | + # Generate py files |
| 59 | + generated_sources = proto_common.declare_generated_files( |
| 60 | + actions = ctx.actions, |
| 61 | + proto_info = proto_info, |
| 62 | + extension = "_pb2.py", |
| 63 | + name_mapper = lambda name: name.replace("-", "_").replace(".", "/"), |
| 64 | + ) |
| 65 | + |
| 66 | + # Handles multiple repository and virtual import cases |
| 67 | + if proto_root.startswith(ctx.bin_dir.path): |
| 68 | + proto_root = proto_root[len(ctx.bin_dir.path) + 1:] |
| 69 | + |
| 70 | + plugin_output = ctx.bin_dir.path + "/" + proto_root |
| 71 | + proto_root = ctx.workspace_name + "/" + proto_root |
| 72 | + |
| 73 | + proto_common.compile( |
| 74 | + actions = ctx.actions, |
| 75 | + proto_info = proto_info, |
| 76 | + proto_lang_toolchain_info = proto_lang_toolchain_info, |
| 77 | + generated_files = generated_sources, |
| 78 | + plugin_output = plugin_output, |
| 79 | + ) |
| 80 | + |
| 81 | + # Generated sources == Python sources |
| 82 | + python_sources = generated_sources |
| 83 | + |
| 84 | + deps = _filter_provider(_PyProtoInfo, getattr(_proto_library, "deps", [])) |
| 85 | + runfiles_from_proto_deps = depset( |
| 86 | + transitive = [dep[DefaultInfo].default_runfiles.files for dep in api_deps] + |
| 87 | + [dep.runfiles_from_proto_deps for dep in deps], |
| 88 | + ) |
| 89 | + transitive_sources = depset( |
| 90 | + direct = python_sources, |
| 91 | + transitive = [dep.transitive_sources for dep in deps], |
| 92 | + ) |
| 93 | + |
| 94 | + return [ |
| 95 | + _PyProtoInfo( |
| 96 | + imports = depset( |
| 97 | + # Adding to PYTHONPATH so the generated modules can be |
| 98 | + # imported. This is necessary when there is |
| 99 | + # strip_import_prefix, the Python modules are generated under |
| 100 | + # _virtual_imports. But it's undesirable otherwise, because it |
| 101 | + # will put the repo root at the top of the PYTHONPATH, ahead of |
| 102 | + # directories added through `imports` attributes. |
| 103 | + [proto_root] if "_virtual_imports" in proto_root else [], |
| 104 | + transitive = [dep[PyInfo].imports for dep in api_deps] + [dep.imports for dep in deps], |
| 105 | + ), |
| 106 | + runfiles_from_proto_deps = runfiles_from_proto_deps, |
| 107 | + transitive_sources = transitive_sources, |
| 108 | + ), |
| 109 | + ] |
| 110 | + |
| 111 | +_py_proto_aspect = aspect( |
| 112 | + implementation = _py_proto_aspect_impl, |
| 113 | + attrs = { |
| 114 | + "_aspect_proto_toolchain": attr.label( |
| 115 | + default = "//python:python_toolchain", |
| 116 | + ), |
| 117 | + }, |
| 118 | + attr_aspects = ["deps"], |
| 119 | + required_providers = [ProtoInfo], |
| 120 | + provides = [_PyProtoInfo], |
| 121 | +) |
| 122 | + |
| 123 | +def _py_proto_library_rule(ctx): |
| 124 | + """Merges results of `py_proto_aspect` in `deps`. |
| 125 | +
|
| 126 | + Args: |
| 127 | + ctx: (RuleContext) The rule context. |
| 128 | + Returns: |
| 129 | + ([PyInfo, DefaultInfo, OutputGroupInfo]) |
| 130 | + """ |
| 131 | + if not ctx.attr.deps: |
| 132 | + fail("'deps' attribute mustn't be empty.") |
| 133 | + |
| 134 | + pyproto_infos = _filter_provider(_PyProtoInfo, ctx.attr.deps) |
| 135 | + default_outputs = depset( |
| 136 | + transitive = [info.transitive_sources for info in pyproto_infos], |
| 137 | + ) |
| 138 | + |
| 139 | + return [ |
| 140 | + DefaultInfo( |
| 141 | + files = default_outputs, |
| 142 | + default_runfiles = ctx.runfiles(transitive_files = depset( |
| 143 | + transitive = |
| 144 | + [default_outputs] + |
| 145 | + [info.runfiles_from_proto_deps for info in pyproto_infos], |
| 146 | + )), |
| 147 | + ), |
| 148 | + OutputGroupInfo( |
| 149 | + default = depset(), |
| 150 | + ), |
| 151 | + PyInfo( |
| 152 | + transitive_sources = default_outputs, |
| 153 | + imports = depset(transitive = [info.imports for info in pyproto_infos]), |
| 154 | + # Proto always produces 2- and 3- compatible source files |
| 155 | + has_py2_only_sources = False, |
| 156 | + has_py3_only_sources = False, |
| 157 | + ), |
| 158 | + ] |
| 159 | + |
| 160 | +py_proto_library = rule( |
| 161 | + implementation = _py_proto_library_rule, |
| 162 | + doc = """ |
| 163 | + Use `py_proto_library` to generate Python libraries from `.proto` files. |
| 164 | +
|
| 165 | + The convention is to name the `py_proto_library` rule `foo_py_pb2`, |
| 166 | + when it is wrapping `proto_library` rule `foo_proto`. |
| 167 | +
|
| 168 | + `deps` must point to a `proto_library` rule. |
| 169 | +
|
| 170 | + Example: |
| 171 | +
|
| 172 | +```starlark |
| 173 | +py_library( |
| 174 | + name = "lib", |
| 175 | + deps = [":foo_py_pb2"], |
| 176 | +) |
| 177 | +
|
| 178 | +py_proto_library( |
| 179 | + name = "foo_py_pb2", |
| 180 | + deps = [":foo_proto"], |
| 181 | +) |
| 182 | +
|
| 183 | +proto_library( |
| 184 | + name = "foo_proto", |
| 185 | + srcs = ["foo.proto"], |
| 186 | +) |
| 187 | +```""", |
| 188 | + attrs = { |
| 189 | + "deps": attr.label_list( |
| 190 | + doc = """ |
| 191 | + The list of `proto_library` rules to generate Python libraries for. |
| 192 | +
|
| 193 | + Usually this is just the one target: the proto library of interest. |
| 194 | + It can be any target providing `ProtoInfo`.""", |
| 195 | + providers = [ProtoInfo], |
| 196 | + aspects = [_py_proto_aspect], |
| 197 | + ), |
| 198 | + }, |
| 199 | + provides = [PyInfo], |
| 200 | +) |
0 commit comments