bazel/python_rules.bzl

   1 """Generates and compiles Python gRPC stubs from proto_library rules."""
   2
   3 load("@grpc_python_dependencies//:requirements.bzl", "requirement")
   4 load(
   5     "//bazel:protobuf.bzl",
   6     "get_include_protoc_args",
   7     "get_plugin_args",
   8     "get_proto_root",
   9     "proto_path_to_generated_filename",
  10 )
  11
  12 _GENERATED_PROTO_FORMAT = "{}_pb2.py"
  13 _GENERATED_GRPC_PROTO_FORMAT = "{}_pb2_grpc.py"
  14
  15 def _get_staged_proto_file(context, source_file):
  16     if source_file.dirname == context.label.package:
  17         return source_file
  18     else:
  19         copied_proto = context.actions.declare_file(source_file.basename)
  20         context.actions.run_shell(
  21             inputs = [source_file],
  22             outputs = [copied_proto],
  23             command = "cp {} {}".format(source_file.path, copied_proto.path),
  24             mnemonic = "CopySourceProto",
  25         )
  26         return copied_proto
  27
  28 def _generate_py_impl(context):
  29     protos = []
  30     for src in context.attr.deps:
  31         for file in src.proto.direct_sources:
  32             protos.append(_get_staged_proto_file(context, file))
  33     includes = [
  34         file
  35         for src in context.attr.deps
  36         for file in src.proto.transitive_imports.to_list()
  37     ]
  38     proto_root = get_proto_root(context.label.workspace_root)
  39     format_str = (_GENERATED_GRPC_PROTO_FORMAT if context.executable.plugin else _GENERATED_PROTO_FORMAT)
  40     out_files = [
  41         context.actions.declare_file(
  42             proto_path_to_generated_filename(
  43                 proto.basename,
  44                 format_str,
  45             ),
  46         )
  47         for proto in protos
  48     ]
  49
  50     arguments = []
  51     tools = [context.executable._protoc]
  52     if context.executable.plugin:
  53         arguments += get_plugin_args(
  54             context.executable.plugin,
  55             context.attr.flags,
  56             context.genfiles_dir.path,
  57             False,
  58         )
  59         tools += [context.executable.plugin]
  60     else:
  61         arguments += [
  62             "--python_out={}:{}".format(
  63                 ",".join(context.attr.flags),
  64                 context.genfiles_dir.path,
  65             ),
  66         ]
  67
  68     arguments += get_include_protoc_args(includes)
  69     arguments += [
  70         "--proto_path={}".format(context.genfiles_dir.path)
  71         for proto in protos
  72     ]
  73     for proto in protos:
  74         massaged_path = proto.path
  75         if massaged_path.startswith(context.genfiles_dir.path):
  76             massaged_path = proto.path[len(context.genfiles_dir.path) + 1:]
  77         arguments.append(massaged_path)
  78
  79     well_known_proto_files = []
  80     if context.attr.well_known_protos:
  81         well_known_proto_directory = context.attr.well_known_protos.files.to_list(
  82         )[0].dirname
  83
  84         arguments += ["-I{}".format(well_known_proto_directory + "/../..")]
  85         well_known_proto_files = context.attr.well_known_protos.files.to_list()
  86
  87     context.actions.run(
  88         inputs = protos + includes + well_known_proto_files,
  89         tools = tools,
  90         outputs = out_files,
  91         executable = context.executable._protoc,
  92         arguments = arguments,
  93         mnemonic = "ProtocInvocation",
  94     )
  95     return struct(files = depset(out_files))
  96
  97 __generate_py = rule(
  98     attrs = {
  99         "deps": attr.label_list(
 100             mandatory = True,
 101             allow_empty = False,
 102             providers = ["proto"],
 103         ),
 104         "plugin": attr.label(
 105             executable = True,
 106             providers = ["files_to_run"],
 107             cfg = "host",
 108         ),
 109         "flags": attr.string_list(
 110             mandatory = False,
 111             allow_empty = True,
 112         ),
 113         "well_known_protos": attr.label(mandatory = False),
 114         "_protoc": attr.label(
 115             default = Label("//external:protocol_compiler"),
 116             executable = True,
 117             cfg = "host",
 118         ),
 119     },
 120     output_to_genfiles = True,
 121     implementation = _generate_py_impl,
 122 )
 123
 124 def _generate_py(well_known_protos, **kwargs):
 125     if well_known_protos:
 126         __generate_py(
 127             well_known_protos = "@com_google_protobuf//:well_known_protos",
 128             **kwargs
 129         )
 130     else:
 131         __generate_py(**kwargs)
 132
 133 def py_proto_library(
 134         name,
 135         deps,
 136         well_known_protos = True,
 137         proto_only = False,
 138         **kwargs):
 139     """Generate python code for a protobuf.
 140
 141     Args:
 142       name: The name of the target.
 143       deps: A list of dependencies. Must contain a single element.
 144       well_known_protos: A bool indicating whether or not to include well-known
 145         protos.
 146       proto_only: A bool indicating whether to generate vanilla protobuf code
 147         or to also generate gRPC code.
 148     """
 149     if len(deps) > 1:
 150         fail("The supported length of 'deps' is 1.")
 151
 152     codegen_target = "_{}_codegen".format(name)
 153     codegen_grpc_target = "_{}_grpc_codegen".format(name)
 154
 155     _generate_py(
 156         name = codegen_target,
 157         deps = deps,
 158         well_known_protos = well_known_protos,
 159         **kwargs
 160     )
 161
 162     if not proto_only:
 163         _generate_py(
 164             name = codegen_grpc_target,
 165             deps = deps,
 166             plugin = "//:grpc_python_plugin",
 167             well_known_protos = well_known_protos,
 168             **kwargs
 169         )
 170
 171         native.py_library(
 172             name = name,
 173             srcs = [
 174                 ":{}".format(codegen_grpc_target),
 175                 ":{}".format(codegen_target),
 176             ],
 177             deps = [requirement("protobuf")],
 178             **kwargs
 179         )
 180     else:
 181         native.py_library(
 182             name = name,
 183             srcs = [":{}".format(codegen_target), ":{}".format(codegen_target)],
 184             deps = [requirement("protobuf")],
 185             **kwargs
 186         )