Refractory of mockserver_textpb_to_binary script to use subprocess.run.

R=harringtond@chromium.org

Bug: 1044242
Change-Id: I05b9eb75f8d1ae507492fbdb987e44719e9b2bca
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2016188Reviewed-by: Dan H <harringtond@chromium.org>
Commit-Queue: Arthur Wang <wuwang@chromium.org>
Cr-Commit-Position: refs/heads/master@{#736113}

components/feed/tools/content_dump.py

@@ -14,15 +14,16 @@
 # Make any desired modifications, and then upload the dump back to the connected
 # device.
 # > content_dump.py --device=FA77D0303076 --apk='com.chrome.canary' --reverse
+import argparse
+import glob
 import os
+import plyvel
+import protoc_util
 import re
-import sys
-import argparse
 import subprocess
-import glob
-from os.path import join, dirname, realpath
+import sys
 
-import plyvel
+from os.path import join, dirname, realpath
 
 # A dynamic import for encoding and decoding of escaped textproto strings.
 _prototext_mod = None
@@ -61,21 +62,8 @@ DUMP_DIR = args.dump_to
 DB_PATH = args.db
 CONTENT_DB_PATH = join(DB_PATH, 'content')
 DEVICE_DB_PATH = "/data/data/{}/app_chrome/Default/feed".format(args.apk)
-_protoc_path = None
-
-
-# Returns the path to the proto compiler, protoc.
-def protoc_path():
-  global _protoc_path
-  if not _protoc_path:
-    protoc_list = list(glob.glob(join(ROOT_DIR, "out") + "/*/protoc")) + list(
-        glob.glob(join(ROOT_DIR, "out") + "/*/*/protoc"))
-    if not len(protoc_list):
-      print("Can't find a suitable build output directory",
-            "(it should have protoc)")
-      sys.exit(1)
-    _protoc_path = protoc_list[0]
-  return _protoc_path
+CONTENT_STORAGE_PROTO = (
+    'components/feed_library/core/proto/content_storage.proto')
 
 
 def adb_base_args():
@@ -97,45 +85,6 @@ def adb_push_db():
                         ["push", CONTENT_DB_PATH, DEVICE_DB_PATH])
 
 
-def get_feed_protos():
-  result = [
-      join(ROOT_DIR, 'components/feed_library/core/proto/content_storage.proto')
-  ]
-  for root, _, files in os.walk(join(ROOT_DIR, "third_party/feed_library")):
-    result += [join(root, f) for f in files if f.endswith('.proto')]
-
-  return result
-
-
-protoc_common_args = [
-    '-I' + join(ROOT_DIR, 'third_party/feed_library/src'), '-I' + join(ROOT_DIR)
-] + get_feed_protos()
-
-
-def run_command(args, input):
-  proc = subprocess.run(
-      args,
-      input=input,
-      stdout=subprocess.PIPE,
-      stderr=subprocess.PIPE,
-      check=True)
-  return proc.stdout
-
-
-# Decode a binary proto into textproto format.
-def decode_proto(data, message_name):
-  return run_command(
-      [protoc_path(), '--decode=' + message_name] + protoc_common_args,
-      data).decode('utf-8')
-
-
-# Encode a textproto into binary proto format.
-def encode_proto(text, message_name):
-  return run_command(
-      [protoc_path(), '--encode=' + message_name] + protoc_common_args,
-      text.encode())
-
-
 # Ignore DB entries with the 'sp::' prefix, as they are not yet supported.
 def is_key_supported(key):
   return not key.startswith('sp::')
@@ -155,13 +104,15 @@ def proto_message_from_db_key(key):
 def extract_db_entry(key, data):
   # DB entries are feed.ContentStorageProto messages. First extract
   # the content_data contained within.
-  text_proto = decode_proto(data, 'feed.ContentStorageProto')
+  text_proto = protoc_util.decode_proto(data, 'feed.ContentStorageProto',
+                                        ROOT_DIR, CONTENT_STORAGE_PROTO)
   m = re.search(r"content_data: \"((?:\\\"|[^\"])*)\"", text_proto)
   raw_data = prototext().CUnescape(m.group(1))
 
   # Next, convert raw_data into a textproto. The DB key informs which message
   # is stored.
-  result = decode_proto(raw_data, proto_message_from_db_key(key))
+  result = protoc_util.decode_proto(raw_data, proto_message_from_db_key(key),
+                                    ROOT_DIR, CONTENT_STORAGE_PROTO)
   return result
 
 
@@ -196,15 +147,17 @@ def load():
         key = file.read().strip()
       with open(join(DUMP_DIR, f), 'r') as file:
         value_text_proto = file.read()
-      value_encoded = encode_proto(value_text_proto,
-                                   proto_message_from_db_key(key))
+      value_encoded = protoc_util.encode_proto(value_text_proto,
+                                               proto_message_from_db_key(key),
+                                               ROOT_DIR, CONTENT_STORAGE_PROTO)
       # Create binary feed.ContentStorageProto by encoding its textproto.
       content_storage_text = 'key: "{}"\ncontent_data: "{}"'.format(
           prototext().CEscape(key, False),
           prototext().CEscape(value_encoded, False))
 
-      store_encoded = encode_proto(content_storage_text,
-                                   'feed.ContentStorageProto')
+      store_encoded = protoc_util.encode_proto(content_storage_text,
+                                               'feed.ContentStorageProto',
+                                               ROOT_DIR, CONTENT_STORAGE_PROTO)
       db.put(key.encode(), store_encoded)
   db.close()
   adb_push_db()

components/feed/tools/mockserver_textpb_to_binary.py

@@ -11,18 +11,19 @@ the proto_convertor script, then a engineer runs this script to encode the
 mockserver textpb file into a binary proto file that is being used by the feed
 card render test (Refers to go/create-a-feed-card-render-test for more).
 
-Make sure you have absl-py installed via pip install absl-py.
+Make sure you have absl-py installed via 'python3 -m pip install absl-py'.
 
 Usage example:
-    python ./mockserver_textpb_to_binary.py
+    python3 ./mockserver_textpb_to_binary.py
     --chromium_path ~/chromium/src
     --output_file /tmp/binary.pb
-    --source_file ~/tmp/original.textpb
+    --source_file /tmp/original.textpb
     --alsologtostderr
 """
 
 import glob
 import os
+import protoc_util
 import subprocess
 
 from absl import app
@@ -35,25 +36,8 @@ flags.DEFINE_string('output_file', '', 'The target output binary file path.')
 flags.DEFINE_string('source_file', '',
                     'The source proto file, in textpb format, path.')
 
-CMD_TEMPLATE = r'cat {} | {} --encode={} -I{} -I$(find {} -name "*.proto") > {}'
 ENCODE_NAMESPACE = 'components.feed.core.proto.wire.mockserver.MockServer'
 COMPONENT_FEED_PROTO_PATH = 'components/feed/core/proto'
-_protoc_path = None
-
-
-def protoc_path(root_dir):
-  """Returns the path to the proto compiler, protoc."""
-  global _protoc_path
-  if not _protoc_path:
-    protoc_list = list(
-        glob.glob(os.path.join(root_dir, "out") + "/*/protoc")) + list(
-            glob.glob(os.path.join(root_dir, "out") + "/*/*/protoc"))
-    if not len(protoc_list):
-      print("Can't find a suitable build output directory",
-            "(it should have protoc)")
-      sys.exit(1)
-    _protoc_path = protoc_list[0]
-  return _protoc_path
 
 
 def main(argv):
@@ -66,12 +50,14 @@ def main(argv):
   if not FLAGS.output_file:
     raise app.UsageError('output_file flag must be set.')
 
-  protoc_cmd = CMD_TEMPLATE.format(
-      FLAGS.source_file, protoc_path(FLAGS.chromium_path), ENCODE_NAMESPACE,
-      FLAGS.chromium_path,
-      os.path.join(FLAGS.chromium_path,
-                   COMPONENT_FEED_PROTO_PATH), FLAGS.output_file)
-  subprocess.call(protoc_cmd, shell=True)
+  with open(FLAGS.source_file) as file:
+    value_text_proto = file.read()
+
+  encoded = protoc_util.encode_proto(value_text_proto, ENCODE_NAMESPACE,
+                                     FLAGS.chromium_path,
+                                     COMPONENT_FEED_PROTO_PATH)
+  with open(FLAGS.output_file, 'wb') as file:
+    file.write(encoded)
 
 
 if __name__ == '__main__':

components/feed/tools/protoc_util.py

+#!/usr/bin/python3
+# 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.
+
+# Lint as: python3
+"""The tools provides lot of protoc related helper functions."""
+
+import glob
+import os
+import subprocess
+
+_protoc_path = None
+
+
+def run_command(args, input):
+  """Uses subprocess to execute the command line args."""
+  proc = subprocess.run(
+      args,
+      input=input,
+      stdout=subprocess.PIPE,
+      stderr=subprocess.PIPE,
+      check=True)
+  return proc.stdout
+
+
+def get_protoc_common_args(root_dir, proto_path):
+  """Returns a list of protoc common args as a list."""
+  result = [
+      '-I' + os.path.join(root_dir, 'third_party/feed_library/src'),
+      '-I' + os.path.join(root_dir)
+  ]
+  for root, _, files in os.walk(os.path.join(root_dir, proto_path)):
+    result += [os.path.join(root, f) for f in files if f.endswith('.proto')]
+
+  return result
+
+
+def encode_proto(text, message_name, root_dir, proto_path):
+  """Calls a command line to encode the text string and returns binary bytes."""
+  return run_command([protoc_path(root_dir), '--encode=' + message_name]
+                     + get_protoc_common_args(root_dir, proto_path),
+                     text.encode())
+
+
+def decode_proto(data, message_name, root_dir, proto_path):
+  """Calls a command line to decode the binary bytes array into text string."""
+  return run_command([protoc_path(root_dir), '--decode=' + message_name
+                     ] + get_protoc_common_args(root_dir, proto_path),
+                     data).decode('utf-8')
+
+
+def protoc_path(root_dir):
+  """Returns the path to the proto compiler, protoc."""
+  global _protoc_path
+  if not _protoc_path:
+    protoc_list = list(
+        glob.glob(os.path.join(root_dir, "out") + "/*/protoc")) + list(
+            glob.glob(os.path.join(root_dir, "out") + "/*/*/protoc"))
+    if not len(protoc_list):
+      print("Can't find a suitable build output directory",
+            "(it should have protoc)")
+      sys.exit(1)
+    _protoc_path = protoc_list[0]
+  return _protoc_path