blob: e1b641e27a73608674c38f8d600a8bc28c698325 [file] [log] [blame]
# -*- 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')
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."""'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]), 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 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 changes.
directory (str): Path to be cleaned up.
"""'Cleaning old files.')
for dirpath, _dirnames, filenames in os.walk(directory):
old = [os.path.join(dirpath, f) for f in filenames if f.endswith('')]
# Remove empty init files to clean up otherwise empty directories.
if '' in filenames:
init = os.path.join(dirpath, '')
if not osutils.ReadFile(init):
for current in old:
def _GenerateFiles(source, output):
"""Generate the proto files from the |source| tree into |output|.
source (str): Path to the proto source root directory.
output (str): Path to the output root directory.
"""'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 =
cmd, cwd=source, print_cmd=False, check=False)
if result.returncode:
raise GenerationError('Error compiling the proto. See the output for a '
def _InstallMissingInits(directory):
"""Add any files not present in the generated protobuf folders."""'Adding missing files.')
for dirpath, _dirnames, filenames in os.walk(directory):
if '' not in filenames:
osutils.Touch(os.path.join(dirpath, ''))
def _PostprocessFiles(directory):
"""Do postprocessing on the generated files.
directory (str): The root directory containing the generated files that are
to be processed.
"""'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('')]
if pb2:
cmd = from_sed + pb2, 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.
output (str|None): The output directory.
source = os.path.join(_PROTO_DIR, 'src')
output = output or os.path.join(_API_DIR, 'gen')
_GenerateFiles(source, output)
def GetParser():
"""Build the argument parser."""
parser = commandline.ArgumentParser(description=__doc__)
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)
return opts
def main(argv):
opts = _ParseArguments(argv)
except Error as e:
return 1