1#!/usr/bin/env python
2# Copyright (c) 2012 The Chromium 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
6"""
7A simple wrapper for protoc.
8
9- Adds includes in generated headers.
10- Handles building with system protobuf as an option.
11"""
12
13import optparse
14import os.path
15import shutil
16import subprocess
17import sys
18import tempfile
19
20PROTOC_INCLUDE_POINT = '// @@protoc_insertion_point(includes)\n'
21
22def ModifyHeader(header_file, extra_header):
23  """Adds |extra_header| to |header_file|. Returns 0 on success.
24
25  |extra_header| is the name of the header file to include.
26  |header_file| is a generated protobuf cpp header.
27  """
28  include_point_found = False
29  header_contents = []
30  with open(header_file) as f:
31    for line in f:
32      header_contents.append(line)
33      if line == PROTOC_INCLUDE_POINT:
34        extra_header_msg = '#include "%s"\n' % extra_header
35        header_contents.append(extra_header_msg)
36        include_point_found = True;
37  if not include_point_found:
38    return 1
39
40  with open(header_file, 'wb') as f:
41    f.write(''.join(header_contents))
42  return 0
43
44
45def RewriteProtoFilesForSystemProtobuf(path):
46  wrapper_dir = tempfile.mkdtemp()
47  try:
48    for filename in os.listdir(path):
49      if not filename.endswith('.proto'):
50        continue
51      with open(os.path.join(path, filename), 'r') as src_file:
52        with open(os.path.join(wrapper_dir, filename), 'w') as dst_file:
53          for line in src_file:
54            # Remove lines that break build with system protobuf.
55            # We cannot optimize for lite runtime, because system lite runtime
56            # does not have a Chromium-specific hack to retain unknown fields.
57            # Similarly, it does not understand corresponding option to control
58            # the usage of that hack.
59            if 'LITE_RUNTIME' in line or 'retain_unknown_fields' in line:
60              continue
61            dst_file.write(line)
62
63    return wrapper_dir
64  except:
65    shutil.rmtree(wrapper_dir)
66    raise
67
68
69def main(argv):
70  parser = optparse.OptionParser()
71  parser.add_option('--include', dest='extra_header',
72                    help='The extra header to include. This must be specified '
73                         'along with --protobuf.')
74  parser.add_option('--protobuf', dest='generated_header',
75                    help='The c++ protobuf header to add the extra header to. '
76                         'This must be specified along with --include.')
77  parser.add_option('--proto-in-dir',
78                    help='The directory containing .proto files.')
79  parser.add_option('--proto-in-file', help='Input file to compile.')
80  parser.add_option('--use-system-protobuf', type=int, default=0,
81                    help='Option to use system-installed protobuf '
82                         'instead of bundled one.')
83  (options, args) = parser.parse_args(sys.argv)
84  if len(args) < 2:
85    return 1
86
87  proto_path = options.proto_in_dir
88  if options.use_system_protobuf == 1:
89    proto_path = RewriteProtoFilesForSystemProtobuf(proto_path)
90  try:
91    # Run what is hopefully protoc.
92    protoc_args = args[1:]
93    protoc_args += ['--proto_path=%s' % proto_path,
94                    os.path.join(proto_path, options.proto_in_file)]
95    ret = subprocess.call(protoc_args)
96    if ret != 0:
97      return ret
98  finally:
99    if options.use_system_protobuf == 1:
100      # Remove temporary directory holding re-written files.
101      shutil.rmtree(proto_path)
102
103  # protoc succeeded, check to see if the generated cpp header needs editing.
104  if not options.extra_header or not options.generated_header:
105    return 0
106  return ModifyHeader(options.generated_header, options.extra_header)
107
108
109if __name__ == '__main__':
110  sys.exit(main(sys.argv))
111