tools/go/mirror_lib.py

"""Provides helper functions for mirroring Go repositories."""

import json
import subprocess
from typing import List, Dict
import unittest.mock

from bazel_tools.tools.python.runfiles import runfiles


def read_file(filepath: str) -> str:
    """Reads an entire file by returning its contents as a string."""
    with open(filepath, "r") as file:
        return file.read()


def parse_go_repositories(filepath: str) -> List[Dict[str, str]]:
    """Parses the top-level go_deps.bzl file.

    This function can parse both the original version of the file generated by
    gazelle as well as the tweaked version generated by
    tweak_gazelle_go_deps.py. The two versions are identical other than what function they call.
    """
    global_functions = {
        "load": unittest.mock.MagicMock(),
        # The gazelle generated version uses go_repository().
        "go_repository": unittest.mock.MagicMock(),
        # The tweak_gazelle_go_deps.py generated version uses
        # maybe_override_go_dep().
        "maybe_override_go_dep": unittest.mock.MagicMock()
    }
    compiled_code = compile(read_file(filepath), filepath, "exec")
    eval(compiled_code, global_functions)

    # Extract the repositories defined in the go_dependencies() function from
    # go_deps.bzl.
    global_functions["go_dependencies"]()

    repositories = []
    for repo_kind in ("go_repository", "maybe_override_go_dep"):
        for repo in global_functions[repo_kind].mock_calls:
            _, _, kwargs = repo
            repositories.append(kwargs)

    return repositories


def parse_go_mirror_info(filepath: str) -> Dict[str, Dict[str, str]]:
    """Parses the tools/go/go_mirrors.bzl file and returns the GO_MIRROR_INFO dictionary."""
    global_data = {}
    compiled_code = compile(read_file(filepath), filepath, "exec")
    eval(compiled_code, global_data)

    return global_data["GO_MIRROR_INFO"]


def write_go_mirror_info(filepath: str, mirror_info: Dict[str, Dict[str,
                                                                    str]]):
    """Saves the specified mirror_info as GO_MIRROR_INFO into tools/go/go_mirrors.bzl."""
    with open(filepath, "w") as file:
        file.write("# This file is auto-generated. Do not edit.\n")
        file.write("GO_MIRROR_INFO = ")
        # Format as JSON first. It's parsable as Starlark.
        json.dump(mirror_info, file, indent=4, sort_keys=True)
        file.write("\n")

    # Properly format the file now so that the linter doesn't complain.
    r = runfiles.Create()
    subprocess.run([
        r.Rlocation(
            "com_github_bazelbuild_buildtools/buildifier/buildifier_/buildifier"
        ),
        filepath,
    ],
                   check=True)