catatester/expand_capabilities.py

"""Script to expand capabilities into products.

Example usage:
  python $0 --capability_list='CAST_AUDIO_RECEIVER,CAST_VIDEO_RECEIVER'
"""
import argparse
import os

from catatester.genfiles import build_config_pb2

from google.protobuf import text_format


BUILD_CONFIG = os.path.join(os.path.dirname(__file__),
                            'configs',
                            'imported',
                            'build_config.protoascii')


def ParseCmdOptions():
  """Parses command arguments."""
  parser = argparse.ArgumentParser()
  parser.add_argument('--capability_list',
                      help='Comma-delimited capabilities',
                      required=True)
  return parser.parse_args()


def GetProductsByCapability(capability):
  """Returns set of products with the given capability."""
  print('Capability name: {}'.format(capability))
  product_set = set()
  try:
    capability_index = (
        build_config_pb2.OtaBuild.Capability.Value(capability))
    print('Capability Index: {}'.format(capability_index))
  except ValueError:
    return product_set

  with open(BUILD_CONFIG) as build:
    builds = text_format.Merge(build.read(), build_config_pb2.BuildConfig())
  assert builds

  for target in builds.ota_build:
    build_target = builds.ota_build[target]
    if build_target.capabilities:
      if capability_index in list(build_target.capabilities):
        product_set.add(build_target.product)

  return product_set


def main():
  assert os.path.isfile(BUILD_CONFIG)

  flags = ParseCmdOptions()

  total_product_set = set()
  for cap in flags.capability_list.split(','):
    total_product_set.update(GetProductsByCapability(cap))

  print('All products: {}'.format(sorted(total_product_set)))


if __name__ == '__main__':
  main()