Add support for kythe annotation in protoc_wrapper

This adds necessary annotation to cc files to allow jumping between
protobuf and cc.

Bug: 1059084
Change-Id: I80ac9f90623acf4fb893a177dafd5840f7e80125
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2134913
Reviewed-by: Scott Violet <sky@chromium.org>
Commit-Queue: Josip Sokcevic <sokcevic@google.com>
Cr-Original-Commit-Position: refs/heads/master@{#757236}
Cr-Mirrored-From: https://chromium.googlesource.com/chromium/src
Cr-Mirrored-Commit: 066775a26c0840d2b43654227092c3e3f1d20ae0
diff --git a/protoc_wrapper.py b/protoc_wrapper.py
index d6e127f..4bb3f06 100755
--- a/protoc_wrapper.py
+++ b/protoc_wrapper.py
@@ -69,10 +69,10 @@
 
 def main(argv):
   parser = argparse.ArgumentParser()
-  parser.add_argument("--protoc",
+  parser.add_argument("--protoc", required=True,
                       help="Relative path to compiler.")
 
-  parser.add_argument("--proto-in-dir",
+  parser.add_argument("--proto-in-dir", required=True,
                       help="Base directory with source protos.")
   parser.add_argument("--cc-out-dir",
                       help="Output directory for standard C++ generator.")
@@ -81,6 +81,9 @@
   parser.add_argument("--plugin-out-dir",
                       help="Output directory for custom generator plugin.")
 
+  parser.add_argument('--enable-kythe-annotations', action='store_true',
+                      help='Enable generation of Kythe kzip, used for '
+                      'codesearch.')
   parser.add_argument("--plugin",
                       help="Relative path to custom generator plugin.")
   parser.add_argument("--plugin-options",
@@ -95,7 +98,7 @@
   parser.add_argument("protos", nargs="+",
                       help="Input protobuf definition file(s).")
 
-  options = parser.parse_args()
+  options = parser.parse_args(argv)
 
   proto_dir = os.path.relpath(options.proto_in_dir)
   protoc_cmd = [os.path.realpath(options.protoc)]
@@ -109,7 +112,19 @@
 
   if options.cc_out_dir:
     cc_out_dir = options.cc_out_dir
-    cc_options = FormatGeneratorOptions(options.cc_options)
+    cc_options_list = []
+    if options.enable_kythe_annotations:
+      cc_options_list.extend([
+          'annotate_headers', 'annotation_pragma_name=kythe_metadata',
+          'annotation_guard_name=KYTHE_IS_RUNNING'
+      ])
+
+    # cc_options will likely have trailing colon so needs to be inserted at the
+    # end.
+    if options.cc_options:
+      cc_options_list.append(options.cc_options)
+
+    cc_options = FormatGeneratorOptions(','.join(cc_options_list))
     protoc_cmd += ["--cpp_out", cc_options + cc_out_dir]
     for filename in protos:
       stripped_name = StripProtoExtension(filename)
@@ -146,7 +161,7 @@
 
 if __name__ == "__main__":
   try:
-    main(sys.argv)
+    main(sys.argv[1:])
   except RuntimeError as e:
     print(e, file=sys.stderr)
     sys.exit(1)
diff --git a/protoc_wrapper_test.py b/protoc_wrapper_test.py
new file mode 100755
index 0000000..ef89a95
--- /dev/null
+++ b/protoc_wrapper_test.py
@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+# Copyright 2020 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+"""Tests for protoc_wrapper."""
+
+from __future__ import print_function
+
+import logging
+import sys
+import unittest
+
+if sys.version_info.major == 2:
+  from StringIO import StringIO
+  import mock
+else:
+  from io import StringIO
+  from unittest import mock
+
+import protoc_wrapper
+
+
+class ProtocWrapperTest(unittest.TestCase):
+  @mock.patch('subprocess.call', return_value=0)
+  def test_minimal_arguments(self, mock_call):
+    protoc_wrapper.main(
+        ['--proto-in-dir', './', '--protoc', '/foo/protoc', 'foo.proto'])
+    mock_call.assert_called_once_with(
+        ['/foo/protoc', '--proto_path', '.', './foo.proto'])
+
+  @mock.patch('subprocess.call', return_value=0)
+  def test_kythe_no_out(self, mock_call):
+    protoc_wrapper.main([
+        '--proto-in-dir', './', '--enable-kythe-annotation', '--protoc',
+        '/foo/protoc', 'foo.proto'
+    ])
+    mock_call.assert_called_once_with(
+        ['/foo/protoc', '--proto_path', '.', './foo.proto'])
+
+  @mock.patch('subprocess.call', return_value=0)
+  def test_kythe_cpp_out_no_options(self, mock_call):
+    protoc_wrapper.main([
+        '--proto-in-dir', './', '--enable-kythe-annotation', '--cc-out-dir',
+        './bar', '--protoc', '/foo/protoc', 'foo.proto'
+    ])
+    mock_call.assert_called_once_with([
+        '/foo/protoc', '--cpp_out',
+        'annotate_headers,annotation_pragma_name=kythe_metadata,annotation_guard_name=KYTHE_IS_RUNNING:./bar',
+        '--proto_path', '.', './foo.proto'
+    ])
+
+  @mock.patch('subprocess.call', return_value=0)
+  def test_kythe_cpp_out_with_options(self, mock_call):
+    protoc_wrapper.main([
+        '--proto-in-dir', './', '--enable-kythe-annotation', '--cc-options',
+        'foo=bar:', '--cc-out-dir', './bar', '--protoc', '/foo/protoc',
+        'foo.proto'
+    ])
+    mock_call.assert_called_once_with([
+        '/foo/protoc', '--cpp_out',
+        'annotate_headers,annotation_pragma_name=kythe_metadata,annotation_guard_name=KYTHE_IS_RUNNING,foo=bar:./bar',
+        '--proto_path', '.', './foo.proto'
+    ])
+
+  @mock.patch('subprocess.call', return_value=0)
+  def test_kythe_cpp_out_with_options_no_colon(self, mock_call):
+    protoc_wrapper.main([
+        '--proto-in-dir', './', '--enable-kythe-annotation', '--cc-options',
+        'foo=bar', '--cc-out-dir', './bar', '--protoc', '/foo/protoc',
+        'foo.proto'
+    ])
+    mock_call.assert_called_once_with([
+        '/foo/protoc', '--cpp_out',
+        'annotate_headers,annotation_pragma_name=kythe_metadata,annotation_guard_name=KYTHE_IS_RUNNING,foo=bar:./bar',
+        '--proto_path', '.', './foo.proto'
+    ])
+
+  @mock.patch('subprocess.call', return_value=0)
+  def test_cpp_out_with_options_no_colon(self, mock_call):
+    protoc_wrapper.main([
+        '--proto-in-dir', './', '--cc-options', 'foo=bar:', '--cc-out-dir',
+        './bar', '--protoc', '/foo/protoc', 'foo.proto'
+    ])
+    mock_call.assert_called_once_with([
+        '/foo/protoc', '--cpp_out', 'foo=bar:./bar', '--proto_path', '.',
+        './foo.proto'
+    ])
+
+
+if __name__ == '__main__':
+  logging.basicConfig(
+      level=logging.DEBUG if '-v' in sys.argv else logging.ERROR)
+  unittest.main()