1# Copyright 2021 The gRPC Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Generates and compiles Python gRPC stubs from proto_library rules."""
15
16load("@rules_proto//proto:defs.bzl", "ProtoInfo")
17load(
18    "//bazel:protobuf.bzl",
19    "declare_out_files",
20    "get_include_directory",
21    "get_out_dir",
22    "get_plugin_args",
23    "get_proto_arguments",
24    "get_staged_proto_file",
25    "includes_from_deps",
26    "is_well_known",
27    "protos_from_context",
28)
29
30_GENERATED_PROTO_FORMAT = "{}_pb2.py"
31_GENERATED_GRPC_PROTO_FORMAT = "{}_pb2_grpc.py"
32
33PyProtoInfo = provider(
34    "The Python outputs from the Protobuf compiler.",
35    fields = {
36        "py_info": "A PyInfo provider for the generated code.",
37        "generated_py_srcs": "The direct (not transitive) generated Python source files.",
38    },
39)
40
41def _merge_pyinfos(pyinfos):
42    return PyInfo(
43        transitive_sources = depset(transitive = [p.transitive_sources for p in pyinfos]),
44        imports = depset(transitive = [p.imports for p in pyinfos]),
45    )
46
47def _gen_py_aspect_impl(target, context):
48    # Early return for well-known protos.
49    if is_well_known(str(context.label)):
50        return [
51            PyProtoInfo(
52                py_info = context.attr._protobuf_library[PyInfo],
53                generated_py_srcs = [],
54            ),
55        ]
56
57    protos = []
58    for p in target[ProtoInfo].direct_sources:
59        protos.append(get_staged_proto_file(target.label, context, p))
60
61    includes = depset(direct = protos, transitive = [target[ProtoInfo].transitive_imports])
62    out_files = declare_out_files(protos, context, _GENERATED_PROTO_FORMAT)
63    generated_py_srcs = out_files
64
65    tools = [context.executable._protoc]
66
67    out_dir = get_out_dir(protos, context)
68
69    arguments = ([
70        "--python_out={}".format(out_dir.path),
71    ] + [
72        "--proto_path={}".format(get_include_directory(i))
73        for i in includes.to_list()
74    ] + [
75        "--proto_path={}".format(context.genfiles_dir.path),
76    ])
77
78    arguments += get_proto_arguments(protos, context.genfiles_dir.path)
79
80    context.actions.run(
81        inputs = protos + includes.to_list(),
82        tools = tools,
83        outputs = out_files,
84        executable = context.executable._protoc,
85        arguments = arguments,
86        mnemonic = "ProtocInvocation",
87    )
88
89    imports = []
90    if out_dir.import_path:
91        imports.append("{}/{}".format(context.workspace_name, out_dir.import_path))
92
93    py_info = PyInfo(transitive_sources = depset(direct = out_files), imports = depset(direct = imports))
94    return PyProtoInfo(
95        py_info = _merge_pyinfos(
96            [
97                py_info,
98                context.attr._protobuf_library[PyInfo],
99            ] + [dep[PyProtoInfo].py_info for dep in context.rule.attr.deps],
100        ),
101        generated_py_srcs = generated_py_srcs,
102    )
103
104_gen_py_aspect = aspect(
105    implementation = _gen_py_aspect_impl,
106    attr_aspects = ["deps"],
107    fragments = ["py"],
108    attrs = {
109        "_protoc": attr.label(
110            default = Label("//external:protocol_compiler"),
111            providers = ["files_to_run"],
112            executable = True,
113            cfg = "host",
114        ),
115        "_protobuf_library": attr.label(
116            default = Label("@com_google_protobuf//:protobuf_python"),
117            providers = [PyInfo],
118        ),
119    },
120)
121
122def _generate_py_impl(context):
123    if (len(context.attr.deps) != 1):
124        fail("Can only compile a single proto at a time.")
125
126    py_sources = []
127
128    # If the proto_library this rule *directly* depends on is in another
129    # package, then we generate .py files to import them in this package. This
130    # behavior is needed to allow rearranging of import paths to make Bazel
131    # outputs align with native python workflows.
132    #
133    # Note that this approach is vulnerable to protoc defining __all__ or other
134    # symbols with __ prefixes that need to be directly imported. Since these
135    # names are likely to be reserved for private APIs, the risk is minimal.
136    if context.label.package != context.attr.deps[0].label.package:
137        for py_src in context.attr.deps[0][PyProtoInfo].generated_py_srcs:
138            reimport_py_file = context.actions.declare_file(py_src.basename)
139            py_sources.append(reimport_py_file)
140            import_line = "from %s import *" % py_src.short_path.replace("/", ".")[:-len(".py")]
141            context.actions.write(reimport_py_file, import_line)
142
143    # Collect output PyInfo provider.
144    imports = [context.label.package + "/" + i for i in context.attr.imports]
145    py_info = PyInfo(transitive_sources = depset(direct = py_sources), imports = depset(direct = imports))
146    out_pyinfo = _merge_pyinfos([py_info, context.attr.deps[0][PyProtoInfo].py_info])
147
148    runfiles = context.runfiles(files = out_pyinfo.transitive_sources.to_list()).merge(context.attr._protobuf_library[DefaultInfo].data_runfiles)
149    return [
150        DefaultInfo(
151            files = out_pyinfo.transitive_sources,
152            runfiles = runfiles,
153        ),
154        out_pyinfo,
155    ]
156
157py_proto_library = rule(
158    attrs = {
159        "deps": attr.label_list(
160            mandatory = True,
161            allow_empty = False,
162            providers = [ProtoInfo],
163            aspects = [_gen_py_aspect],
164        ),
165        "_protoc": attr.label(
166            default = Label("//external:protocol_compiler"),
167            providers = ["files_to_run"],
168            executable = True,
169            cfg = "host",
170        ),
171        "_protobuf_library": attr.label(
172            default = Label("@com_google_protobuf//:protobuf_python"),
173            providers = [PyInfo],
174        ),
175        "imports": attr.string_list(),
176    },
177    implementation = _generate_py_impl,
178)
179
180def _generate_pb2_grpc_src_impl(context):
181    protos = protos_from_context(context)
182    includes = includes_from_deps(context.attr.deps)
183    out_files = declare_out_files(protos, context, _GENERATED_GRPC_PROTO_FORMAT)
184
185    plugin_flags = ["grpc_2_0"] + context.attr.strip_prefixes
186
187    arguments = []
188    tools = [context.executable._protoc, context.executable._grpc_plugin]
189    out_dir = get_out_dir(protos, context)
190    arguments += get_plugin_args(
191        context.executable._grpc_plugin,
192        plugin_flags,
193        out_dir.path,
194        False,
195    )
196
197    arguments += [
198        "--proto_path={}".format(get_include_directory(i))
199        for i in includes
200    ]
201    arguments.append("--proto_path={}".format(context.genfiles_dir.path))
202    arguments += get_proto_arguments(protos, context.genfiles_dir.path)
203
204    context.actions.run(
205        inputs = protos + includes,
206        tools = tools,
207        outputs = out_files,
208        executable = context.executable._protoc,
209        arguments = arguments,
210        mnemonic = "ProtocInvocation",
211    )
212
213    p = PyInfo(transitive_sources = depset(direct = out_files))
214    py_info = _merge_pyinfos(
215        [
216            p,
217            context.attr._grpc_library[PyInfo],
218        ] + [dep[PyInfo] for dep in context.attr.py_deps],
219    )
220
221    runfiles = context.runfiles(files = out_files, transitive_files = py_info.transitive_sources).merge(context.attr._grpc_library[DefaultInfo].data_runfiles)
222
223    return [
224        DefaultInfo(
225            files = depset(direct = out_files),
226            runfiles = runfiles,
227        ),
228        py_info,
229    ]
230
231_generate_pb2_grpc_src = rule(
232    attrs = {
233        "deps": attr.label_list(
234            mandatory = True,
235            allow_empty = False,
236            providers = [ProtoInfo],
237        ),
238        "py_deps": attr.label_list(
239            mandatory = True,
240            allow_empty = False,
241            providers = [PyInfo],
242        ),
243        "strip_prefixes": attr.string_list(),
244        "_grpc_plugin": attr.label(
245            executable = True,
246            providers = ["files_to_run"],
247            cfg = "host",
248            default = Label("//src/compiler:grpc_python_plugin"),
249        ),
250        "_protoc": attr.label(
251            executable = True,
252            providers = ["files_to_run"],
253            cfg = "host",
254            default = Label("//external:protocol_compiler"),
255        ),
256        "_grpc_library": attr.label(
257            default = Label("//src/python/grpcio/grpc:grpcio"),
258            providers = [PyInfo],
259        ),
260    },
261    implementation = _generate_pb2_grpc_src_impl,
262)
263
264def py_grpc_library(
265        name,
266        srcs,
267        deps,
268        strip_prefixes = [],
269        **kwargs):
270    """Generate python code for gRPC services defined in a protobuf.
271
272    Args:
273      name: The name of the target.
274      srcs: (List of `labels`) a single proto_library target containing the
275        schema of the service.
276      deps: (List of `labels`) a single py_proto_library target for the
277        proto_library in `srcs`.
278      strip_prefixes: (List of `strings`) If provided, this prefix will be
279        stripped from the beginning of foo_pb2 modules imported by the
280        generated stubs. This is useful in combination with the `imports`
281        attribute of the `py_library` rule.
282      **kwargs: Additional arguments to be supplied to the invocation of
283        py_library.
284    """
285    if len(srcs) != 1:
286        fail("Can only compile a single proto at a time.")
287
288    if len(deps) != 1:
289        fail("Deps must have length 1.")
290
291    _generate_pb2_grpc_src(
292        name = name,
293        deps = srcs,
294        py_deps = deps,
295        strip_prefixes = strip_prefixes,
296        **kwargs
297    )
298