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"""Utility functions for generating protobuf code."""
15
16load("@rules_proto//proto:defs.bzl", "ProtoInfo")
17
18_PROTO_EXTENSION = ".proto"
19_VIRTUAL_IMPORTS = "/_virtual_imports/"
20
21_WELL_KNOWN_PROTOS_BASE = [
22    "any_proto",
23    "api_proto",
24    "compiler_plugin_proto",
25    "descriptor_proto",
26    "duration_proto",
27    "empty_proto",
28    "field_mask_proto",
29    "source_context_proto",
30    "struct_proto",
31    "timestamp_proto",
32    "type_proto",
33    "wrappers_proto",
34]
35
36def well_known_proto_libs():
37    return ["@com_google_protobuf//:" + b for b in _WELL_KNOWN_PROTOS_BASE]
38
39def is_well_known(label):
40    # Bazel surfaces labels as their undelying identity, even if they are referenced
41    # via aliases. Bazel also does not currently provide a way to find the real label
42    # underlying an alias. So the implementation detail that the WKTs present at the
43    # top level of the protobuf repo are actually backed by targets in the
44    # //src/google/protobuf package leaks through here.
45    # We include both the alias path and the underlying path to be resilient to
46    # reversions of this change as well as for continuing compatiblity with repos
47    # that happen to pull in older versions of protobuf.
48    all_wkt_targets = (["@com_google_protobuf//:" + b for b in _WELL_KNOWN_PROTOS_BASE] +
49                       ["@com_google_protobuf//src/google/protobuf:" + b for b in _WELL_KNOWN_PROTOS_BASE])
50    return label in all_wkt_targets
51
52def get_proto_root(workspace_root):
53    """Gets the root protobuf directory.
54
55    Args:
56      workspace_root: context.label.workspace_root
57
58    Returns:
59      The directory relative to which generated include paths should be.
60    """
61    if workspace_root:
62        return "/{}".format(workspace_root)
63    else:
64        return ""
65
66def _strip_proto_extension(proto_filename):
67    if not proto_filename.endswith(_PROTO_EXTENSION):
68        fail('"{}" does not end with "{}"'.format(
69            proto_filename,
70            _PROTO_EXTENSION,
71        ))
72    return proto_filename[:-len(_PROTO_EXTENSION)]
73
74def proto_path_to_generated_filename(proto_path, fmt_str):
75    """Calculates the name of a generated file for a protobuf path.
76
77    For example, "examples/protos/helloworld.proto" might map to
78      "helloworld.pb.h".
79
80    Args:
81      proto_path: The path to the .proto file.
82      fmt_str: A format string used to calculate the generated filename. For
83        example, "{}.pb.h" might be used to calculate a C++ header filename.
84
85    Returns:
86      The generated filename.
87    """
88    return fmt_str.format(_strip_proto_extension(proto_path))
89
90def get_include_directory(source_file):
91    """Returns the include directory path for the source_file.
92
93    All of the include statements within the given source_file are calculated
94    relative to the directory returned by this method.
95
96    The returned directory path can be used as the "--proto_path=" argument
97    value.
98
99    Args:
100      source_file: A proto file.
101
102    Returns:
103      The include directory path for the source_file.
104    """
105    directory = source_file.path
106    prefix_len = 0
107
108    if is_in_virtual_imports(source_file):
109        root, relative = source_file.path.split(_VIRTUAL_IMPORTS, 2)
110        result = root + _VIRTUAL_IMPORTS + relative.split("/", 1)[0]
111        return result
112
113    if not source_file.is_source and directory.startswith(source_file.root.path):
114        prefix_len = len(source_file.root.path) + 1
115
116    if directory.startswith("external", prefix_len):
117        external_separator = directory.find("/", prefix_len)
118        repository_separator = directory.find("/", external_separator + 1)
119        return directory[:repository_separator]
120    else:
121        return source_file.root.path if source_file.root.path else "."
122
123def get_plugin_args(
124        plugin,
125        flags,
126        dir_out,
127        generate_mocks,
128        plugin_name = "PLUGIN"):
129    """Returns arguments configuring protoc to use a plugin for a language.
130
131    Args:
132      plugin: An executable file to run as the protoc plugin.
133      flags: The plugin flags to be passed to protoc.
134      dir_out: The output directory for the plugin.
135      generate_mocks: A bool indicating whether to generate mocks.
136      plugin_name: A name of the plugin, it is required to be unique when there
137      are more than one plugin used in a single protoc command.
138    Returns:
139      A list of protoc arguments configuring the plugin.
140    """
141    augmented_flags = list(flags)
142    if generate_mocks:
143        augmented_flags.append("generate_mock_code=true")
144
145    augmented_dir_out = dir_out
146    if augmented_flags:
147        augmented_dir_out = ",".join(augmented_flags) + ":" + dir_out
148
149    return [
150        "--plugin=protoc-gen-{plugin_name}={plugin_path}".format(
151            plugin_name = plugin_name,
152            plugin_path = plugin.path,
153        ),
154        "--{plugin_name}_out={dir_out}".format(
155            plugin_name = plugin_name,
156            dir_out = augmented_dir_out,
157        ),
158    ]
159
160def _make_prefix(label):
161    """Returns the directory prefix for a label.
162
163    @repo//foo/bar:sub/dir/file.proto  =>  'external/repo/foo/bar/'
164    //foo/bar:sub/dir/file.proto       =>  'foo/bar/'
165    //:sub/dir/file.proto              =>  ''
166
167    That is, the prefix can be removed from a file's full path to
168    obtain the file's relative location within the package's effective
169    directory."""
170
171    wsr = label.workspace_root
172    pkg = label.package
173
174    if not wsr and not pkg:
175        return ""
176    elif not wsr:
177        return pkg + "/"
178    elif not pkg:
179        return wsr + "/"
180    else:
181        return wsr + "/" + pkg + "/"
182
183def get_staged_proto_file(label, context, source_file):
184    """Copies a proto file to the appropriate location if necessary.
185
186    Args:
187      label: The label of the rule using the .proto file.
188      context: The ctx object for the rule or aspect.
189      source_file: The original .proto file.
190
191    Returns:
192      The original proto file OR a new file in the staged location.
193    """
194    if source_file.dirname == label.package or \
195       is_in_virtual_imports(source_file):
196        # Current target and source_file are in same package
197        return source_file
198    else:
199        # Current target and source_file are in different packages (most
200        # probably even in different repositories)
201        prefix = _make_prefix(source_file.owner)
202        copied_proto = context.actions.declare_file(source_file.path[len(prefix):])
203        context.actions.run_shell(
204            inputs = [source_file],
205            outputs = [copied_proto],
206            command = "cp {} {}".format(source_file.path, copied_proto.path),
207            mnemonic = "CopySourceProto",
208        )
209        return copied_proto
210
211def protos_from_context(context):
212    """Copies proto files to the appropriate location.
213
214    Args:
215      context: The ctx object for the rule.
216
217    Returns:
218      A list of the protos.
219    """
220    protos = []
221    for src in context.attr.deps:
222        for file in src[ProtoInfo].direct_sources:
223            protos.append(get_staged_proto_file(context.label, context, file))
224    return protos
225
226def includes_from_deps(deps):
227    """Get includes from rule dependencies."""
228    return [
229        file
230        for src in deps
231        for file in src[ProtoInfo].transitive_imports.to_list()
232    ]
233
234def get_proto_arguments(protos, genfiles_dir_path):
235    """Get the protoc arguments specifying which protos to compile.
236
237    Args:
238      protos: The protob files to supply.
239      genfiles_dir_path: The path to the genfiles directory.
240
241    Returns:
242      The arguments to supply to protoc.
243    """
244    arguments = []
245    for proto in protos:
246        strip_prefix_len = 0
247        if is_in_virtual_imports(proto):
248            incl_directory = get_include_directory(proto)
249            if proto.path.startswith(incl_directory):
250                strip_prefix_len = len(incl_directory) + 1
251        elif proto.path.startswith(genfiles_dir_path):
252            strip_prefix_len = len(genfiles_dir_path) + 1
253
254        arguments.append(proto.path[strip_prefix_len:])
255
256    return arguments
257
258def declare_out_files(protos, context, generated_file_format):
259    """Declares and returns the files to be generated.
260
261    Args:
262      protos: A list of files. The protos to declare.
263      context: The context object.
264      generated_file_format: A format string. Will be passed to
265        proto_path_to_generated_filename to generate the filename of each
266        generated file.
267
268    Returns:
269      A list of file providers.
270    """
271
272    out_file_paths = []
273    for proto in protos:
274        if not is_in_virtual_imports(proto):
275            prefix = _make_prefix(proto.owner)
276            full_prefix = context.genfiles_dir.path + "/" + prefix
277            if proto.path.startswith(full_prefix):
278                out_file_paths.append(proto.path[len(full_prefix):])
279            elif proto.path.startswith(prefix):
280                out_file_paths.append(proto.path[len(prefix):])
281        else:
282            out_file_paths.append(proto.path[proto.path.index(_VIRTUAL_IMPORTS) + 1:])
283
284    return [
285        context.actions.declare_file(
286            proto_path_to_generated_filename(
287                out_file_path,
288                generated_file_format,
289            ),
290        )
291        for out_file_path in out_file_paths
292    ]
293
294def get_out_dir(protos, context):
295    """Returns the value to supply to the --<lang>_out= protoc flag.
296
297    The result is based on the input source proto files and current context.
298
299    Args:
300        protos: A list of protos to be used as source files in protoc command
301        context: A ctx object for the rule.
302    Returns:
303        The value of --<lang>_out= argument.
304    """
305    at_least_one_virtual = 0
306    for proto in protos:
307        if is_in_virtual_imports(proto):
308            at_least_one_virtual = True
309        elif at_least_one_virtual:
310            fail("Proto sources must be either all virtual imports or all real")
311    if at_least_one_virtual:
312        out_dir = get_include_directory(protos[0])
313        ws_root = protos[0].owner.workspace_root
314        prefix = "/" + _make_prefix(protos[0].owner) + _VIRTUAL_IMPORTS[1:]
315
316        return struct(
317            path = out_dir,
318            import_path = out_dir[out_dir.find(prefix) + 1:],
319        )
320
321    out_dir = context.genfiles_dir.path
322    ws_root = context.label.workspace_root
323    if ws_root:
324        out_dir = out_dir + "/" + ws_root
325    return struct(path = out_dir, import_path = None)
326
327def is_in_virtual_imports(source_file, virtual_folder = _VIRTUAL_IMPORTS):
328    """Determines if source_file is virtual.
329
330    A file is virtual if placed in the _virtual_imports subdirectory. The
331    output of all proto_library targets which use import_prefix and/or
332    strip_import_prefix arguments is placed under _virtual_imports directory.
333
334    Args:
335        source_file: A proto file.
336        virtual_folder: The virtual folder name (is set to "_virtual_imports"
337            by default)
338    Returns:
339        True if source_file is located under _virtual_imports, False otherwise.
340    """
341    return not source_file.is_source and virtual_folder in source_file.path
342