Home | History | Annotate | Download | only in protoc_wrapper 
      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 """
      7 A simple wrapper for protoc.
      8 
      9 - Adds includes in generated headers.
     10 - Handles building with system protobuf as an option.
     11 """
     12 
     13 import optparse
     14 import os.path
     15 import shutil
     16 import subprocess
     17 import sys
     18 import tempfile
     19 
     20 PROTOC_INCLUDE_POINT = '// @@protoc_insertion_point(includes)\n'
     21 
     22 def 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 
     45 def 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 
     69 def 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 
    109 if __name__ == '__main__':
    110   sys.exit(main(sys.argv))
    111