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