| # Copyright 2018 The ChromiumOS Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Compile the Build API's proto. |
| |
| Install proto using CIPD to ensure a consistent protoc version. |
| """ |
| |
| import enum |
| import logging |
| from pathlib import Path |
| import tempfile |
| from typing import Iterable, Optional |
| |
| from chromite.lib import cipd |
| from chromite.lib import commandline |
| from chromite.lib import constants |
| from chromite.lib import cros_build_lib |
| from chromite.lib import git |
| from chromite.lib import osutils |
| |
| |
| # Chromite's protobuf library version (third_party/google/protobuf). |
| PROTOC_VERSION = "21.9" |
| |
| # Protobuf dropped the major version number after 3.20, jumping to 21.0. But |
| # some places (e.g., in protobuf/__init__.py) refer to this as 4.21.0. |
| PROTOC_MAJOR_VERSION = "4" |
| |
| _CIPD_PACKAGE = "infra/3pp/tools/protoc/linux-amd64" |
| _CIPD_PACKAGE_VERSION = f"version:2@{PROTOC_VERSION}" |
| _CHROMEOS_CONFIG_PATH = constants.SOURCE_ROOT / "src" / "config" |
| |
| |
| class Error(Exception): |
| """Base error class for the module.""" |
| |
| |
| class GenerationError(Error): |
| """A failure we can't recover from.""" |
| |
| |
| @enum.unique |
| class ProtocVersion(enum.Enum): |
| """Enum for possible protoc versions.""" |
| |
| # The SDK version of the bindings use the protoc in the SDK, and so is |
| # compatible with the protobuf library in the SDK, i.e. the one installed |
| # via the ebuild. |
| SDK = enum.auto() |
| # The Chromite version of the bindings uses a protoc binary downloaded from |
| # CIPD that matches the version of the protobuf library in |
| # chromite/third_party/google/protobuf. |
| CHROMITE = enum.auto() |
| # Type annotations for the chromite bindings. |
| CHROMITE_PYI = enum.auto() |
| |
| def get_gen_dir(self) -> Path: |
| """Get the chromite/api directory path.""" |
| if self is ProtocVersion.SDK: |
| return constants.CHROMITE_DIR / "api" / "gen_sdk" |
| return constants.CHROMITE_DIR / "api" / "gen" |
| |
| def get_proto_dir(self) -> Path: |
| """Get the proto directory for the target protoc.""" |
| return constants.SOURCE_ROOT / "infra" / "proto_branched" |
| |
| def get_protoc_command(self, cipd_root: Optional[Path] = None) -> Path: |
| """Get protoc command path.""" |
| if self is ProtocVersion.SDK: |
| return Path("protoc") |
| assert cipd_root |
| return cipd_root / "bin" / "protoc" |
| |
| def get_suffix(self) -> str: |
| """Get the file suffix of generated output.""" |
| return "pyi" if self is ProtocVersion.CHROMITE_PYI else "py" |
| |
| |
| @enum.unique |
| class SubdirectorySet(enum.Enum): |
| """Enum for the subsets of the proto to compile.""" |
| |
| ALL = enum.auto() |
| DEFAULT = enum.auto() |
| |
| def get_source_dirs( |
| self, source: Path, chromeos_config_path: Path |
| ) -> Iterable[Path]: |
| """Get the directories for the given subdirectory set.""" |
| if self is self.ALL: |
| return [ |
| source, |
| chromeos_config_path / "proto" / "chromiumos", |
| ] |
| |
| subdirs = [ |
| source / "analysis_service", |
| source / "chromite", |
| source / "chromiumos", |
| source / "config", |
| source / "test_platform", |
| source / "device", |
| source / "ide_query", |
| chromeos_config_path / "proto" / "chromiumos", |
| ] |
| return subdirs |
| |
| |
| def check_upstream_changes(repo: Path) -> None: |
| """Ensures the checked-out branch at `repo` includes the upstream HEAD.""" |
| branch = git.GetCurrentBranch(repo) |
| if branch: |
| upstream = git.GetTrackingBranchViaGitConfig( |
| repo, branch, for_checkout=False |
| ) |
| if not upstream: |
| cros_build_lib.Die("Failed to get upstream for %s.", repo) |
| else: |
| # Detached head. |
| branch = "HEAD" |
| upstream = git.RemoteRef("cros", "refs/heads/main") |
| remote_head = git.RunGit( |
| repo, |
| ["ls-remote", upstream.remote, upstream.ref], |
| ).stdout.split(maxsplit=1)[0] |
| logging.notice( |
| "Ensuring proto dir %s contains %s from %s", |
| repo.relative_to(constants.SOURCE_ROOT), |
| upstream.ref, |
| upstream.remote, |
| ) |
| branches = git.RunGit( |
| repo, |
| ["branch", "--contains", remote_head, branch], |
| print_cmd=True, |
| check=False, |
| ) |
| # Git emits the branch name if the commit is in the history. Otherwise, the |
| # commit is either not found (git exits with an error - 129), or missing |
| # from the branch. In either case, there is no output to stdout. |
| if not branches.stdout: |
| cros_build_lib.Die( |
| "The checked-out branch (%s) at %s is missing the remote head.\n" |
| "Please repo sync (and rebase), or skip this check by passing" |
| " --no-check-upstream-proto-changes-included.", |
| branch, |
| repo, |
| ) |
| |
| |
| def InstallProtoc(protoc_version: ProtocVersion) -> Path: |
| """Install protoc from CIPD.""" |
| if protoc_version is ProtocVersion.SDK: |
| cipd_root = None |
| else: |
| cipd_root = Path( |
| cipd.InstallPackage( |
| cipd.GetCIPDFromCache(), _CIPD_PACKAGE, _CIPD_PACKAGE_VERSION |
| ) |
| ) |
| return protoc_version.get_protoc_command(cipd_root) |
| |
| |
| def _CleanTargetDirectory( |
| directory: Path, protoc_version: ProtocVersion |
| ) -> None: |
| """Remove any existing generated files in the directory. |
| |
| This clean only removes the generated files to avoid accidentally destroying |
| __init__.py customizations down the line. That will leave otherwise empty |
| directories in place if things get moved. Neither case is relevant at the |
| time of writing, but lingering empty directories seemed better than |
| diagnosing accidental __init__.py changes. |
| |
| Args: |
| directory: Path to be cleaned up. |
| protoc_version: The type of generated files to be cleaned. |
| """ |
| logging.info("Cleaning old files from %s.", directory) |
| for current in directory.rglob(f"*_pb2.{protoc_version.get_suffix()}"): |
| # Remove old generated files. |
| current.unlink() |
| |
| # Note the generator does not currently make __init__.pyi files but, if it |
| # did, we'd want them to be cleaned up here. |
| for current in directory.rglob(f"__init__.{protoc_version.get_suffix()}"): |
| # Remove empty init files to clean up otherwise empty directories. |
| if not current.stat().st_size: |
| current.unlink() |
| |
| |
| def _GenerateFiles( |
| source: Path, |
| output: Path, |
| protoc_version: ProtocVersion, |
| dir_subset: SubdirectorySet, |
| protoc_bin_path: Path, |
| ) -> None: |
| """Generate the proto files from the |source| tree into |output|. |
| |
| Args: |
| source: Path to the proto source root directory. |
| output: Path to the output root directory. |
| protoc_version: Which protoc to use. |
| dir_subset: The subset of the proto to compile. |
| protoc_bin_path: The protoc command to use. |
| """ |
| logging.info("Generating %s files to %s.", protoc_version, output) |
| osutils.SafeMakedirs(output) |
| |
| targets = [] |
| |
| chromeos_config_path = _CHROMEOS_CONFIG_PATH |
| |
| with tempfile.TemporaryDirectory() as tempdir: |
| if not chromeos_config_path.exists(): |
| chromeos_config_path = Path(tempdir) / "config" |
| |
| logging.info("Creating shallow clone of chromiumos/config") |
| git.Clone( |
| chromeos_config_path, |
| "%s/chromiumos/config" % constants.EXTERNAL_GOB_URL, |
| depth=1, |
| ) |
| |
| for src_dir in dir_subset.get_source_dirs(source, chromeos_config_path): |
| targets.extend(list(src_dir.rglob("*.proto"))) |
| |
| output_type = ( |
| "pyi" if protoc_version is ProtocVersion.CHROMITE_PYI else "python" |
| ) |
| cmd = [ |
| protoc_bin_path, |
| "-I", |
| chromeos_config_path / "proto", |
| f"--{output_type}_out", |
| output, |
| "--proto_path", |
| source, |
| ] |
| cmd.extend(targets) |
| |
| result = cros_build_lib.dbg_run( |
| cmd, |
| cwd=source, |
| check=False, |
| enter_chroot=protoc_version is ProtocVersion.SDK, |
| ) |
| |
| if result.returncode: |
| raise GenerationError( |
| "Error compiling the proto. See the output for a " "message." |
| ) |
| |
| |
| def _InstallMissingInits( |
| directory: Path, protoc_version: ProtocVersion |
| ) -> None: |
| """Add missing __init__.py files in the generated protobuf folders.""" |
| if protoc_version is ProtocVersion.CHROMITE_PYI: |
| # For pyi, rely on module markers left behind by CHROMITE flows to avoid |
| # additional clutter. |
| return |
| |
| logging.info("Adding missing __init__.py files in %s.", directory) |
| # glob ** returns only directories. |
| for current in directory.rglob("**"): |
| (current / "__init__.py").touch() |
| |
| |
| def _PostprocessFiles(directory: Path, protoc_version: ProtocVersion) -> None: |
| """Do postprocessing on the generated files. |
| |
| Args: |
| directory: The root directory containing the generated files that are |
| to be processed. |
| protoc_version: Which protoc is being used to generate the files. |
| """ |
| logging.info("Postprocessing: Fix imports in %s.", directory) |
| # We are using a negative address here (the /address/! portion of the sed |
| # command) to make sure we don't change any imports from protobuf itself. |
| address = "^from google.protobuf" |
| # Find: 'from x import y_pb2 as x_dot_y_pb2'. |
| # "\(^google.protobuf[^ ]*\)" matches the module we're importing from. |
| # - \( and \) are for groups in sed. |
| # - ^google.protobuf prevents changing the import for protobuf's files. |
| # - [^ ] = Not a space. The [:space:] character set is too broad, but |
| # would technically work too. |
| find = r"^from \([^ ]*\) import \([^ ]*\)_pb2 as \([^ ]*\)$" |
| # Substitute: 'from chromite.api.gen[_sdk].x import y_pb2 as x_dot_y_pb2'. |
| if protoc_version is ProtocVersion.SDK: |
| sub = "from chromite.api.gen_sdk.\\1 import \\2_pb2 as \\3" |
| else: |
| sub = "from chromite.api.gen.\\1 import \\2_pb2 as \\3" |
| |
| seds = [f"/{address}/!s/{find}/{sub}/g"] |
| |
| if protoc_version is ProtocVersion.CHROMITE: |
| # We also need to change the google.protobuf imports to point directly |
| # at the chromite.third_party version of the library. |
| # The SDK version of the proto is meant to be used with the protobuf |
| # libraries installed in the SDK, so leave those as google.protobuf. |
| # For CHROMITE_PYI, types for the protobuf imports come from type stubs, |
| # which won't map if the imports are renamed to third_party. |
| g_p_address = "^from google.protobuf" |
| g_p_find = r"from \([^ ]*\) import \(.*\)$" |
| g_p_sub = "from chromite.third_party.\\1 import \\2" |
| seds.append(f"/{g_p_address}/s/{g_p_find}/{g_p_sub}/") |
| |
| if protoc_version is ProtocVersion.CHROMITE_PYI: |
| # Workaround https://github.com/protocolbuffers/protobuf/issues/11402 |
| # until cl/560557754 in in the protoc version used. |
| untyped_empty_slot_list = "s/\\(^ *__slots__ = \\)\\[]/\\1()/" |
| |
| # Suppress errors on GRPC options field numbers using ClassVar outside |
| # of a class body. See b/297782342. |
| ignore_class_var_in_file_scope = ( |
| "s/^[A-Z_]*: _ClassVar\\[int]/& # type: ignore/" |
| ) |
| seds.extend((untyped_empty_slot_list, ignore_class_var_in_file_scope)) |
| |
| pb2 = list(directory.rglob(f"*_pb2.{protoc_version.get_suffix()}")) |
| if pb2: |
| cros_build_lib.dbg_run(["sed", "-i", ";".join(seds)] + pb2) |
| |
| |
| def CompileProto( |
| protoc_version: ProtocVersion, |
| output: Optional[Path] = None, |
| dir_subset: SubdirectorySet = SubdirectorySet.DEFAULT, |
| postprocess: bool = True, |
| ) -> None: |
| """Compile the Build API protobuf files. |
| |
| By default, this will compile from infra/proto/src to api/gen. The output |
| directory may be changed, but the imports will always be treated as if it is |
| in the default location. |
| |
| Args: |
| output: The output directory. |
| protoc_version: Which protoc to use for the compilation. |
| dir_subset: What proto to compile. |
| postprocess: Whether to run the postprocess step. |
| """ |
| source = protoc_version.get_proto_dir() / "src" |
| if not output: |
| output = protoc_version.get_gen_dir() |
| |
| protoc_bin_path = InstallProtoc(protoc_version) |
| _CleanTargetDirectory(output, protoc_version) |
| _GenerateFiles(source, output, protoc_version, dir_subset, protoc_bin_path) |
| _InstallMissingInits(output, protoc_version) |
| if postprocess: |
| _PostprocessFiles(output, protoc_version) |
| |
| |
| def GetParser(): |
| """Build the argument parser.""" |
| parser = commandline.ArgumentParser(description=__doc__) |
| parser.add_bool_argument( |
| "--check-upstream-proto-changes-included", |
| True, |
| "Check whether the checked-out proto branch includes changes" |
| " from refs/heads/main on the remote.", |
| "Skip the check that validates whether upstream proto changes" |
| " are included in the current branch.", |
| ) |
| standard_group = parser.add_argument_group( |
| "Committed Bindings", |
| description="Options for generating the bindings in chromite/api/.", |
| ) |
| standard_group.add_argument( |
| "--chromite", |
| dest="protoc_version", |
| action="append_const", |
| const=ProtocVersion.CHROMITE, |
| help="Generate only the chromite bindings. Generates all by default. " |
| "The chromite bindings are compatible with the version of protobuf " |
| "in chromite/third_party.", |
| ) |
| standard_group.add_argument( |
| "--pyi", |
| dest="protoc_version", |
| action="append_const", |
| const=ProtocVersion.CHROMITE_PYI, |
| help="Generate only the pyi type annotations.", |
| ) |
| standard_group.add_argument( |
| "--sdk", |
| dest="protoc_version", |
| action="append_const", |
| const=ProtocVersion.SDK, |
| help="Generate only the SDK bindings. Generates all by default. The " |
| "SDK bindings are compiled by protoc in the SDK, and is compatible " |
| "with the version of protobuf in the SDK (i.e. the one installed by " |
| "the ebuild).", |
| ) |
| |
| dest_group = parser.add_argument_group( |
| "Out of Tree Bindings", |
| description="Options for generating bindings in a custom location.", |
| ) |
| dest_group.add_argument( |
| "--destination", |
| type="str_path", |
| help="A directory where a single version of the proto should be " |
| "generated. When not given, the proto generates in all default " |
| "locations instead.", |
| ) |
| dest_group.add_argument( |
| "--dest-sdk", |
| action="store_const", |
| dest="dest_protoc", |
| default=ProtocVersion.CHROMITE, |
| const=ProtocVersion.SDK, |
| help="Generate the SDK version of the protos in --destination instead " |
| "of the chromite version.", |
| ) |
| dest_group.add_argument( |
| "--all-proto", |
| action="store_const", |
| dest="dir_subset", |
| default=SubdirectorySet.DEFAULT, |
| const=SubdirectorySet.ALL, |
| help="Compile ALL proto instead of just the subset needed for the API. " |
| "Only considered when generating out of tree bindings.", |
| ) |
| dest_group.add_argument( |
| "--skip-postprocessing", |
| action="store_false", |
| dest="postprocess", |
| default=True, |
| help="Skip postprocessing files.", |
| ) |
| return parser |
| |
| |
| def _ParseArguments(argv): |
| """Parse and validate arguments.""" |
| parser = GetParser() |
| opts = parser.parse_args(argv) |
| |
| if not opts.protoc_version: |
| opts.protoc_version = [ |
| ProtocVersion.CHROMITE, |
| ProtocVersion.SDK, |
| ProtocVersion.CHROMITE_PYI, |
| ] |
| |
| if opts.destination: |
| opts.destination = Path(opts.destination) |
| |
| opts.Freeze() |
| return opts |
| |
| |
| def main(argv): |
| opts = _ParseArguments(argv) |
| |
| if opts.destination: |
| # Destination set, only compile a single version in the destination. |
| try: |
| CompileProto( |
| protoc_version=opts.dest_protoc, |
| output=opts.destination, |
| dir_subset=opts.dir_subset, |
| postprocess=opts.postprocess, |
| ) |
| except Error as e: |
| cros_build_lib.Die("Error compiling bindings to destination: %s", e) |
| else: |
| return 0 |
| |
| if ProtocVersion.CHROMITE in opts.protoc_version: |
| # Validate here to avoid checking again inside the chroot. |
| if opts.check_upstream_proto_changes_included: |
| check_upstream_changes(ProtocVersion.CHROMITE.get_proto_dir()) |
| if _CHROMEOS_CONFIG_PATH.exists(): |
| check_upstream_changes(_CHROMEOS_CONFIG_PATH) |
| |
| # Compile the chromite bindings. |
| try: |
| CompileProto(protoc_version=ProtocVersion.CHROMITE) |
| except Error as e: |
| cros_build_lib.Die("Error compiling chromite bindings: %s", e) |
| |
| if ProtocVersion.CHROMITE_PYI in opts.protoc_version: |
| try: |
| CompileProto(protoc_version=ProtocVersion.CHROMITE_PYI) |
| except Error as e: |
| cros_build_lib.Die("Error compiling type annotations: %s", e) |
| |
| if ProtocVersion.SDK in opts.protoc_version: |
| # Compile the SDK bindings. |
| if not cros_build_lib.IsInsideChroot(): |
| # Rerun inside the SDK instead of trying to map all the paths. |
| cmd = [ |
| ( |
| constants.CHROOT_SOURCE_ROOT |
| / "chromite" |
| / "api" |
| / "compile_build_api_proto" |
| ), |
| "--sdk", |
| ] |
| result = cros_build_lib.dbg_run(cmd, enter_chroot=True, check=False) |
| return result.returncode |
| else: |
| try: |
| CompileProto(protoc_version=ProtocVersion.SDK) |
| except Error as e: |
| cros_build_lib.Die("Error compiling SDK bindings: %s", e) |