blob: 4b0f8754da55d795bebd2c55b7d656b2fc44166a [file] [log] [blame]
Alex Kleinf4dc4f52018-12-05 13:55:12 -07001# -*- coding: utf-8 -*-
2# Copyright 2018 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
Alex Kleinf9859972019-03-14 17:11:42 -06006"""Compile the Build API's proto.
7
8Install proto using CIPD to ensure a consistent protoc version.
9"""
Alex Kleinf4dc4f52018-12-05 13:55:12 -070010
11from __future__ import print_function
12
13import os
14
15from chromite.lib import commandline
Alex Kleinc33c1912019-02-15 10:29:13 -070016from chromite.lib import constants
Alex Kleinf4dc4f52018-12-05 13:55:12 -070017from chromite.lib import cros_build_lib
Alex Klein5534f992019-09-16 16:31:23 -060018from chromite.lib import cros_logging as logging
Alex Kleinf9859972019-03-14 17:11:42 -060019from chromite.lib import osutils
20
21_API_DIR = os.path.join(constants.CHROMITE_DIR, 'api')
22_CIPD_ROOT = os.path.join(constants.CHROMITE_DIR, '.cipd_bin')
23_PROTOC = os.path.join(_CIPD_ROOT, 'protoc')
24_PROTO_DIR = os.path.join(constants.CHROMITE_DIR, 'infra', 'proto')
25
26PROTOC_VERSION = '3.6.1'
27
28
Alex Klein5534f992019-09-16 16:31:23 -060029class Error(Exception):
30 """Base error class for the module."""
31
32
33class GenerationError(Error):
34 """A failure we can't recover from."""
35
36
Alex Kleinf9859972019-03-14 17:11:42 -060037def _InstallProtoc():
38 """Install protoc from CIPD."""
Alex Klein5534f992019-09-16 16:31:23 -060039 logging.info('Installing protoc.')
Alex Kleinf9859972019-03-14 17:11:42 -060040 cmd = ['cipd', 'ensure']
41 # Clean up the output.
42 cmd.extend(['-log-level', 'warning'])
43 # Set the install location.
44 cmd.extend(['-root', _CIPD_ROOT])
45
46 ensure_content = ('infra/tools/protoc/${platform} '
47 'protobuf_version:v%s' % PROTOC_VERSION)
48 with osutils.TempDir() as tempdir:
49 ensure_file = os.path.join(tempdir, 'cipd_ensure_file')
50 osutils.WriteFile(ensure_file, ensure_content)
51
52 cmd.extend(['-ensure-file', ensure_file])
53
Alex Klein5534f992019-09-16 16:31:23 -060054 cros_build_lib.RunCommand(cmd, cwd=constants.CHROMITE_DIR, print_cmd=False)
55
Alex Kleinf9859972019-03-14 17:11:42 -060056
57def _CleanTargetDirectory(directory):
58 """Remove any existing generated files in the directory.
59
60 This clean only removes the generated files to avoid accidentally destroying
61 __init__.py customizations down the line. That will leave otherwise empty
62 directories in place if things get moved. Neither case is relevant at the
63 time of writing, but lingering empty directories seemed better than
64 diagnosing accidental __init__.py changes.
65
66 Args:
67 directory (str): Path to be cleaned up.
68 """
Alex Klein5534f992019-09-16 16:31:23 -060069 logging.info('Cleaning old files.')
Alex Kleinf9859972019-03-14 17:11:42 -060070 for dirpath, _dirnames, filenames in os.walk(directory):
71 old = [os.path.join(dirpath, f) for f in filenames if f.endswith('_pb2.py')]
Alex Klein5534f992019-09-16 16:31:23 -060072 # Remove empty init files to clean up otherwise empty directories.
73 if '__init__.py' in filenames:
74 init = os.path.join(dirpath, '__init__.py')
75 if not osutils.ReadFile(init):
76 old.append(init)
77
Alex Kleinf9859972019-03-14 17:11:42 -060078 for current in old:
79 osutils.SafeUnlink(current)
80
Alex Klein5534f992019-09-16 16:31:23 -060081
Alex Kleinf9859972019-03-14 17:11:42 -060082def _GenerateFiles(source, output):
83 """Generate the proto files from the |source| tree into |output|.
84
85 Args:
86 source (str): Path to the proto source root directory.
87 output (str): Path to the output root directory.
88 """
Alex Klein5534f992019-09-16 16:31:23 -060089 logging.info('Generating files.')
Alex Kleinf9859972019-03-14 17:11:42 -060090 targets = []
91
92 # Only compile the subset we need for the API.
Alex Klein5534f992019-09-16 16:31:23 -060093 subdirs = [
94 os.path.join(source, 'chromite'),
95 os.path.join(source, 'chromiumos'),
96 os.path.join(source, 'test_platform')
97 ]
Alex Kleinf9859972019-03-14 17:11:42 -060098 for basedir in subdirs:
99 for dirpath, _dirnames, filenames in os.walk(basedir):
100 for filename in filenames:
101 if filename.endswith('.proto'):
102 # We have a match, add the file.
103 targets.append(os.path.join(dirpath, filename))
104
Alex Klein5534f992019-09-16 16:31:23 -0600105 cmd = [_PROTOC, '--python_out', output, '--proto_path', source] + targets
106 result = cros_build_lib.RunCommand(
107 cmd, cwd=source, print_cmd=False, error_code_ok=True)
108
109 if result.returncode:
110 raise GenerationError('Error compiling the proto. See the output for a '
111 'message.')
Alex Kleinf9859972019-03-14 17:11:42 -0600112
113
114def _InstallMissingInits(directory):
115 """Add any __init__.py files not present in the generated protobuf folders."""
Alex Klein5534f992019-09-16 16:31:23 -0600116 logging.info('Adding missing __init__.py files.')
Alex Kleinf9859972019-03-14 17:11:42 -0600117 for dirpath, _dirnames, filenames in os.walk(directory):
118 if '__init__.py' not in filenames:
119 osutils.Touch(os.path.join(dirpath, '__init__.py'))
120
121
122def _PostprocessFiles(directory):
123 """Do postprocessing on the generated files.
124
125 Args:
126 directory (str): The root directory containing the generated files that are
127 to be processed.
128 """
Alex Klein5534f992019-09-16 16:31:23 -0600129 logging.info('Postprocessing: Fix imports.')
Alex Kleinf9859972019-03-14 17:11:42 -0600130 # We are using a negative address here (the /address/! portion of the sed
131 # command) to make sure we don't change any imports from protobuf itself.
132 address = '^from google.protobuf'
133 # Find: 'from x import y_pb2 as x_dot_y_pb2'.
134 # "\(^google.protobuf[^ ]*\)" matches the module we're importing from.
135 # - \( and \) are for groups in sed.
136 # - ^google.protobuf prevents changing the import for protobuf's files.
137 # - [^ ] = Not a space. The [:space:] character set is too broad, but would
138 # technically work too.
139 find = r'^from \([^ ]*\) import \([^ ]*\)_pb2 as \([^ ]*\)$'
140 # Substitute: 'from chromite.api.gen.x import y_pb2 as x_dot_y_pb2'.
141 sub = 'from chromite.api.gen.\\1 import \\2_pb2 as \\3'
Alex Klein5534f992019-09-16 16:31:23 -0600142 from_sed = [
143 'sed', '-i',
144 '/%(address)s/!s/%(find)s/%(sub)s/g' % {
145 'address': address,
146 'find': find,
147 'sub': sub
148 }
149 ]
Alex Kleinf9859972019-03-14 17:11:42 -0600150
151 for dirpath, _dirnames, filenames in os.walk(directory):
152 # Update the
153 pb2 = [os.path.join(dirpath, f) for f in filenames if f.endswith('_pb2.py')]
154 if pb2:
155 cmd = from_sed + pb2
Alex Klein5534f992019-09-16 16:31:23 -0600156 cros_build_lib.RunCommand(cmd, print_cmd=False)
Alex Kleinf9859972019-03-14 17:11:42 -0600157
158
159def CompileProto(output=None):
160 """Compile the Build API protobuf files.
161
162 By default this will compile from infra/proto/src to api/gen. The output
163 directory may be changed, but the imports will always be treated as if it is
164 in the default location.
165
166 Args:
167 output (str|None): The output directory.
168 """
169 source = os.path.join(_PROTO_DIR, 'src')
170 output = output or os.path.join(_API_DIR, 'gen')
171
172 _InstallProtoc()
173 _CleanTargetDirectory(output)
174 _GenerateFiles(source, output)
175 _InstallMissingInits(output)
176 _PostprocessFiles(output)
Alex Kleinf4dc4f52018-12-05 13:55:12 -0700177
178
179def GetParser():
Alex Kleinf9859972019-03-14 17:11:42 -0600180 """Build the argument parser."""
Alex Kleinf4dc4f52018-12-05 13:55:12 -0700181 parser = commandline.ArgumentParser(description=__doc__)
Alex Klein5534f992019-09-16 16:31:23 -0600182 parser.add_argument(
183 '--destination',
184 type='path',
185 help='The directory where the proto should be generated. Defaults to '
186 'the correct directory for the API.')
Alex Kleinf4dc4f52018-12-05 13:55:12 -0700187 return parser
188
189
190def _ParseArguments(argv):
191 """Parse and validate arguments."""
192 parser = GetParser()
193 opts = parser.parse_args(argv)
194
195 opts.Freeze()
196 return opts
197
198
199def main(argv):
Alex Kleinf9859972019-03-14 17:11:42 -0600200 opts = _ParseArguments(argv)
Alex Kleinf4dc4f52018-12-05 13:55:12 -0700201
Alex Klein5534f992019-09-16 16:31:23 -0600202 try:
203 CompileProto(output=opts.destination)
204 except Error as e:
205 logging.error(e.message)
206 return 1