lib/cros_collections.py

# -*- coding: utf-8 -*-
# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Chromite extensions on top of the collections module."""

from __future__ import print_function


def _CollectionExec(expr, classname):
  """Hack to workaround <=Python-2.7.8 exec bug.

  See https://bugs.python.org/issue21591 for details.

  TODO(crbug.com/998624): Drop this in Jan 2020.
  """
  namespace = {}
  exec(expr, {}, namespace)  # pylint: disable=exec-used
  return namespace[classname]


# We have nested kwargs below, so disable the |kwargs| naming here.
# pylint: disable=docstring-misnamed-args
def Collection(classname, **default_kwargs):
  """Create a new class with mutable named members.

  This is like collections.namedtuple, but mutable.  Also similar to the
  python 3.3 types.SimpleNamespace.

  Examples:
    # Declare default values for this new class.
    Foo = cros_build_lib.Collection('Foo', a=0, b=10)
    # Create a new class but set b to 4.
    foo = Foo(b=4)
    # Print out a (will be the default 0) and b (will be 4).
    print('a = %i, b = %i' % (foo.a, foo.b))
  """

  def sn_init(self, **kwargs):
    """The new class's __init__ function."""
    # First verify the kwargs don't have excess settings.
    valid_keys = set(self.__slots__)
    these_keys = set(kwargs.keys())
    invalid_keys = these_keys - valid_keys
    if invalid_keys:
      raise TypeError('invalid keyword arguments for this object: %r' %
                      invalid_keys)

    # Now initialize this object.
    for k in valid_keys:
      setattr(self, k, kwargs.get(k, default_kwargs[k]))

  def sn_repr(self):
    """The new class's __repr__ function."""
    return '%s(%s)' % (classname, ', '.join(
        '%s=%r' % (k, getattr(self, k)) for k in self.__slots__))

  # Give the new class a unique name and then generate the code for it.
  classname = 'Collection_%s' % classname
  expr = '\n'.join((
      'class %(classname)s(object):',
      '  __slots__ = ["%(slots)s"]',
  )) % {
      'classname': classname,
      'slots': '", "'.join(sorted(default_kwargs)),
  }

  # Create the class in a local namespace as exec requires.
  new_class = _CollectionExec(expr, classname)

  # Bind the helpers.
  new_class.__init__ = sn_init
  new_class.__repr__ = sn_repr

  return new_class
# pylint: enable=docstring-misnamed-args


def GroupByKey(input_iter, key):
  """Split an iterable of dicts, based on value of a key.

  GroupByKey([{'a': 1}, {'a': 2}, {'a': 1, 'b': 2}], 'a') =>
    {1: [{'a': 1}, {'a': 1, 'b': 2}], 2: [{'a': 2}]}

  Args:
    input_iter: An iterable of dicts.
    key: A string specifying the key name to split by.

  Returns:
    A dictionary, mapping from each unique value for |key| that
    was encountered in |input_iter| to a list of entries that had
    that value.
  """
  split_dict = dict()
  for entry in input_iter:
    split_dict.setdefault(entry.get(key), []).append(entry)
  return split_dict


def GroupNamedtuplesByKey(input_iter, key):
  """Split an iterable of namedtuples, based on value of a key.

  Args:
    input_iter: An iterable of namedtuples.
    key: A string specifying the key name to split by.

  Returns:
    A dictionary, mapping from each unique value for |key| that
    was encountered in |input_iter| to a list of entries that had
    that value.
  """
  split_dict = {}
  for entry in input_iter:
    split_dict.setdefault(getattr(entry, key, None), []).append(entry)
  return split_dict


def InvertDictionary(origin_dict):
  """Invert the key value mapping in the origin_dict.

  Given an origin_dict {'key1': {'val1', 'val2'}, 'key2': {'val1', 'val3'},
  'key3': {'val3'}}, the returned inverted dict will be
  {'val1': {'key1', 'key2'}, 'val2': {'key1'}, 'val3': {'key2', 'key3'}}

  Args:
    origin_dict: A dict mapping each key to a group (collection) of values.

  Returns:
    An inverted dict mapping each key to a set of its values.
  """
  new_dict = {}
  for origin_key, origin_values in origin_dict.items():
    for origin_value in origin_values:
      new_dict.setdefault(origin_value, set()).add(origin_key)

  return new_dict