| # -*- coding: utf-8 -*- |
| # Copyright 2018 The Chromium OS Authors. All rights reserved. |
| # 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. |
| """ |
| |
| from __future__ import print_function |
| |
| import os |
| import sys |
| |
| from chromite.lib import commandline |
| from chromite.lib import constants |
| from chromite.lib import cros_build_lib |
| from chromite.lib import cros_logging as logging |
| from chromite.lib import osutils |
| |
| |
| assert sys.version_info >= (3, 6), 'This module requires Python 3.6+' |
| |
| |
| _API_DIR = os.path.join(constants.CHROMITE_DIR, 'api') |
| _CIPD_ROOT = os.path.join(constants.CHROMITE_DIR, '.cipd_bin') |
| _PROTOC = os.path.join(_CIPD_ROOT, 'protoc') |
| _PROTO_DIR = os.path.join(constants.CHROMITE_DIR, 'infra', 'proto') |
| |
| PROTOC_VERSION = '3.6.1' |
| |
| |
| class Error(Exception): |
| """Base error class for the module.""" |
| |
| |
| class GenerationError(Error): |
| """A failure we can't recover from.""" |
| |
| |
| def _InstallProtoc(): |
| """Install protoc from CIPD.""" |
| logging.info('Installing protoc.') |
| cmd = ['cipd', 'ensure'] |
| # Clean up the output. |
| cmd.extend(['-log-level', 'warning']) |
| # Set the install location. |
| cmd.extend(['-root', _CIPD_ROOT]) |
| |
| ensure_content = ('infra/tools/protoc/${platform} ' |
| 'protobuf_version:v%s' % PROTOC_VERSION) |
| with osutils.TempDir() as tempdir: |
| ensure_file = os.path.join(tempdir, 'cipd_ensure_file') |
| osutils.WriteFile(ensure_file, ensure_content) |
| |
| cmd.extend(['-ensure-file', ensure_file]) |
| |
| cros_build_lib.run(cmd, cwd=constants.CHROMITE_DIR, print_cmd=False) |
| |
| |
| def _CleanTargetDirectory(directory): |
| """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 (str): Path to be cleaned up. |
| """ |
| logging.info('Cleaning old files.') |
| for dirpath, _dirnames, filenames in os.walk(directory): |
| old = [os.path.join(dirpath, f) for f in filenames if f.endswith('_pb2.py')] |
| # Remove empty init files to clean up otherwise empty directories. |
| if '__init__.py' in filenames: |
| init = os.path.join(dirpath, '__init__.py') |
| if not osutils.ReadFile(init): |
| old.append(init) |
| |
| for current in old: |
| osutils.SafeUnlink(current) |
| |
| |
| def _GenerateFiles(source, output): |
| """Generate the proto files from the |source| tree into |output|. |
| |
| Args: |
| source (str): Path to the proto source root directory. |
| output (str): Path to the output root directory. |
| """ |
| logging.info('Generating files.') |
| targets = [] |
| |
| # Only compile the subset we need for the API. |
| subdirs = [ |
| os.path.join(source, 'chromite'), |
| os.path.join(source, 'chromiumos'), |
| os.path.join(source, 'client'), |
| os.path.join(source, 'config'), |
| os.path.join(source, 'test_platform'), |
| os.path.join(source, 'device') |
| ] |
| for basedir in subdirs: |
| for dirpath, _dirnames, filenames in os.walk(basedir): |
| for filename in filenames: |
| if filename.endswith('.proto'): |
| # We have a match, add the file. |
| targets.append(os.path.join(dirpath, filename)) |
| |
| cmd = [_PROTOC, '--python_out', output, '--proto_path', source] + targets |
| result = cros_build_lib.run( |
| cmd, cwd=source, print_cmd=False, check=False) |
| |
| if result.returncode: |
| raise GenerationError('Error compiling the proto. See the output for a ' |
| 'message.') |
| |
| |
| def _InstallMissingInits(directory): |
| """Add any __init__.py files not present in the generated protobuf folders.""" |
| logging.info('Adding missing __init__.py files.') |
| for dirpath, _dirnames, filenames in os.walk(directory): |
| if '__init__.py' not in filenames: |
| osutils.Touch(os.path.join(dirpath, '__init__.py')) |
| |
| |
| def _PostprocessFiles(directory): |
| """Do postprocessing on the generated files. |
| |
| Args: |
| directory (str): The root directory containing the generated files that are |
| to be processed. |
| """ |
| logging.info('Postprocessing: Fix imports.') |
| # 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.x import y_pb2 as x_dot_y_pb2'. |
| sub = 'from chromite.api.gen.\\1 import \\2_pb2 as \\3' |
| from_sed = [ |
| 'sed', '-i', |
| '/%(address)s/!s/%(find)s/%(sub)s/g' % { |
| 'address': address, |
| 'find': find, |
| 'sub': sub |
| } |
| ] |
| |
| for dirpath, _dirnames, filenames in os.walk(directory): |
| # Update the |
| pb2 = [os.path.join(dirpath, f) for f in filenames if f.endswith('_pb2.py')] |
| if pb2: |
| cmd = from_sed + pb2 |
| cros_build_lib.run(cmd, print_cmd=False) |
| |
| |
| def CompileProto(output=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 (str|None): The output directory. |
| """ |
| source = os.path.join(_PROTO_DIR, 'src') |
| output = output or os.path.join(_API_DIR, 'gen') |
| |
| _InstallProtoc() |
| _CleanTargetDirectory(output) |
| _GenerateFiles(source, output) |
| _InstallMissingInits(output) |
| _PostprocessFiles(output) |
| |
| |
| def GetParser(): |
| """Build the argument parser.""" |
| parser = commandline.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| '--destination', |
| type='path', |
| help='The directory where the proto should be generated. Defaults to ' |
| 'the correct directory for the API.') |
| return parser |
| |
| |
| def _ParseArguments(argv): |
| """Parse and validate arguments.""" |
| parser = GetParser() |
| opts = parser.parse_args(argv) |
| |
| opts.Freeze() |
| return opts |
| |
| |
| def main(argv): |
| opts = _ParseArguments(argv) |
| |
| try: |
| CompileProto(output=opts.destination) |
| except Error as e: |
| logging.error(e) |
| return 1 |