proto.build_defs

"""Build rules for compiling protocol buffers & gRPC service stubs.

Note that these are some of the most complex of our built-in build rules,
because of their cross-language nature. Each proto_library rule declares a set of
sub-rules to run protoc & the appropriate java_library, go_library rules etc. Users
shouldn't worry about those sub-rules and just declare a dependency directly on
the proto_library rule to get its appropriate outputs.

It is possible to add extra languages to these for generation. This is accomplished
via the 'languages' argument; this can be simply a list of languages to build, but
can also be a mapping of language name -> definition of how to build it. The definition
should be the return value of proto_language.
"""

subinclude(
    "//build_defs:proto_languages",
    "//build_defs:grpc_languages",
)

_DEFAULT_GRPC_LABELS = ['grpc']


def merge_languages(lhs_dict, rhs_dict):
    """Merges dictionary a into dictionary b, overwriting where a's values are not None."""
    if not lhs_dict:
        return rhs_dict

    if not isinstance(lhs_dict, dict):
        return {lang: rhs_dict[lang] for lang in lhs_dict}  # Languages can be passed as just a list.
    return {k: v or rhs_dict[k] for k, v in lhs_dict.items()}

def proto_library(name:str, srcs:list, deps:list=[], visibility:list=None, labels:list&features&tags=[],
                  languages:list|dict=None, test_only:bool&testonly=False, root_dir:str='', protoc_flags:list=[],
                  additional_context:dict=None):
    """Compile a .proto file to generated code for various languages.

    Args:
      name (str): Name of the rule
      srcs (list): Input .proto files.
      deps (list): Dependencies
      visibility (list): Visibility specification for the rule.
      labels (list): List of labels to apply to this rule.
      languages (list | dict): List of languages to generate rules for, chosen from the set {cc, py, go, java, js}.
                               Alternatively, a dict mapping the language name to a definition of how to build it
                               (see proto_language for more details of the values).
      test_only (bool): If True, can only be used in test rules.
      root_dir (str): The directory that the protos are compiled relative to. Useful if your
                      proto files have import statements that are not relative to the repo root.
      protoc_flags (list): Additional flags to pass to protoc. Note that these are inherited by
                           further rules that depend on this one (because in nearly all cases that
                           will be necessary for them to build too).
      additional_context (dict): Additional language context to be passed to the `proto_language()` build definition.
                                 This can be used to provide language specific parameters.
    """
    languages = merge_languages(languages, proto_languages())
    lang_plugins = sorted(languages.items())

    # If the protoc tool is a build rule, depend on it as it's probably a protoc_binary() which provides the well known
    # types
    if CONFIG.PROTO.PROTOC_TOOL.startswith("//") or CONFIG.PROTO.PROTOC_TOOL.startswith(":"):
        # Extract the wkt into another rule as we already depend on the main rule as a tool.
        deps += [filegroup(
            name = name,
            tag = "protoc_wkt",
            exported_deps = [CONFIG.PROTO.PROTOC_TOOL],
            test_only = test_only,
            requires = ["proto"],
        )]


    provides = {}
    for language, plugin in lang_plugins:
        provides[language] = plugin['build_def'](
            name = name,
            language = language,
            srcs = srcs,
            deps = deps,
            visibility = visibility,
            labels = labels,
            test_only = test_only,
            root_dir = root_dir,
            protoc_flags = protoc_flags,
            additional_context = additional_context,
        )
        additional_provides = plugin["additional_provides"]
        if isinstance(additional_provides, dict):
            for k, v in additional_provides.items():
                provides[k] = f":_{name}#{v}"
        else:
            for k in additional_provides:
                provides[k] = f":_{name}#{k}"

    deps =  provides.values()
    provides["proto"] = f":{name}"
    # This is the final rule that directs dependencies to the appropriate language.
    return filegroup(
        name = name,
        srcs = srcs,
        deps = deps,
        provides = provides,
        visibility = visibility,
        labels = labels,
        test_only = test_only,
    )

def grpc_library(name:str, srcs:list, deps:list=None, visibility:list=None, languages:list|dict=None,
                 labels:list&features&tags=[], test_only:bool&testonly=False, root_dir:str='', protoc_flags:list=None):
    """Defines a rule for a grpc library.

    Args:
      name (str): Name of the rule
      srcs (list): Input .proto files.
      deps (list): Dependencies (other grpc_library or proto_library rules)
      visibility (list): Visibility specification for the rule.
      languages (list | dict): List of languages to generate rules for, chosen from the set {cc, py, go, java}.
                               Alternatively, a dict mapping the language name to a definition of how to build it
                               (see proto_language for more details of the values).
      labels (list): List of labels to apply to this rule.
      test_only (bool): If True, this rule can only be used by test rules.
      root_dir (str): The directory that the protos are compiled relative to. Useful if your
                      proto files have import statements that are not relative to the repo root.
      protoc_flags (list): Additional flags to pass to protoc.
    """
    return proto_library(
        name = name,
        srcs = srcs,
        deps = deps,
        languages = merge_languages(languages, grpc_languages()),
        visibility = visibility,
        labels = labels + _DEFAULT_GRPC_LABELS,
        test_only = test_only,
        root_dir = root_dir,
        protoc_flags = protoc_flags,
    )

def protoc_binary(name, url:str|dict = '', version:str = '', hashes=None, deps=None, visibility=None):
    """Downloads a precompiled protoc binary.

    You will obviously need to choose a version that is available on Github - there aren't
    necessarily protoc downloads for every protobuf release.

    Args:
      name (str): Name of the rule
      url (str | dict): The URL used to download protoc. Can be a single string or a dictionary mapping
                        HOSTOS-HOSTARCH to URLs i.e. linux-amd64: 'https://...'. Either provide url or version, but not both.
      version (str): The version of protoc to download (e.g. '3.4.0'). Protoc will be downloaded from https://github.com/protocolbuffers/protobuf/releases/downaload/...
                     and the rule will use the current platforms OS and ARCH setting. Either provide url or version,
                     but not both.
      hashes (list): Hashes to verify the download against.
      deps (list): Any other dependencies
      visibility (list): Visibility of the rule.
    """
    if url and version:
        fail("Either version or url should be provided but not both")

    if CONFIG.HOSTOS == 'darwin':
        HOSTOS='osx'
    else:
        HOSTOS=CONFIG.HOSTOS

    if CONFIG.HOSTARCH == 'amd64':
        HOSTARCH = 'x86_64'
    elif CONFIG.HOSTARCH == 'x86':
        HOSTARCH = 'x86_32'
    elif CONFIG.HOSTARCH == 'arm64':
        if CONFIG.HOSTOS == 'darwin':
            # There's no arm64 distribution, but we can use the x86 version thanks to rosetta!
            HOSTARCH = 'x86_64'
        else:
            HOSTARCH = 'aarch_64'
    else:
        HOSTARCH = CONFIG.HOSTARCH

    if version:
        protoc_url = f'https://github.com/protocolbuffers/protobuf/releases/download/v{version}/protoc-{version}-{HOSTOS}-{HOSTARCH}.zip'
    else:
        protoc_url = url if isinstance(url, str) else url[f'{HOSTOS}-{HOSTARCH}']

    download = remote_file(
        name = name,
        _tag = 'download',
        url = protoc_url,
        hashes = hashes,
        deps = deps,
    )
    # The well-known types are included in the same zipfile.
    wkt = build_rule(
        name = name,
        tag = 'wkt',
        srcs = [download],
        outs = ['google'],
        tools = [CONFIG.JARCAT_TOOL],
        cmd = '$TOOL x $SRCS --strip_prefix include',
        labels = [
            'protoc:-I"$TMP_DIR"/' + package_name(),
        ],
    )
    return genrule(
        name = name,
        srcs = [download],
        outs = ['protoc'],
        tools = [CONFIG.JARCAT_TOOL],
        binary = True,
        cmd = '$TOOL x $SRCS bin/protoc',
        visibility = visibility,
        provides = {'proto': wkt},
    )