1# -*- coding: utf-8 -*- 2# Copyright 2018 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6"""Chromite extensions on top of the collections module.""" 7 8from __future__ import print_function 9 10 11def _CollectionExec(expr, classname): 12 """Hack to workaround <=Python-2.7.8 exec bug. 13 14 See https://bugs.python.org/issue21591 for details. 15 16 TODO(crbug.com/998624): Drop this in Jan 2020. 17 """ 18 namespace = {} 19 exec(expr, {}, namespace) # pylint: disable=exec-used 20 return namespace[classname] 21 22 23# We have nested kwargs below, so disable the |kwargs| naming here. 24# pylint: disable=docstring-misnamed-args 25def Collection(classname, **default_kwargs): 26 """Create a new class with mutable named members. 27 28 This is like collections.namedtuple, but mutable. Also similar to the 29 python 3.3 types.SimpleNamespace. 30 31 Examples: 32 # Declare default values for this new class. 33 Foo = cros_build_lib.Collection('Foo', a=0, b=10) 34 # Create a new class but set b to 4. 35 foo = Foo(b=4) 36 # Print out a (will be the default 0) and b (will be 4). 37 print('a = %i, b = %i' % (foo.a, foo.b)) 38 """ 39 40 def sn_init(self, **kwargs): 41 """The new class's __init__ function.""" 42 # First verify the kwargs don't have excess settings. 43 valid_keys = set(self.__slots__) 44 these_keys = set(kwargs.keys()) 45 invalid_keys = these_keys - valid_keys 46 if invalid_keys: 47 raise TypeError('invalid keyword arguments for this object: %r' % 48 invalid_keys) 49 50 # Now initialize this object. 51 for k in valid_keys: 52 setattr(self, k, kwargs.get(k, default_kwargs[k])) 53 54 def sn_repr(self): 55 """The new class's __repr__ function.""" 56 return '%s(%s)' % (classname, ', '.join( 57 '%s=%r' % (k, getattr(self, k)) for k in self.__slots__)) 58 59 # Give the new class a unique name and then generate the code for it. 60 classname = 'Collection_%s' % classname 61 expr = '\n'.join(( 62 'class %(classname)s(object):', 63 ' __slots__ = ["%(slots)s"]', 64 )) % { 65 'classname': classname, 66 'slots': '", "'.join(sorted(default_kwargs)), 67 } 68 69 # Create the class in a local namespace as exec requires. 70 new_class = _CollectionExec(expr, classname) 71 72 # Bind the helpers. 73 new_class.__init__ = sn_init 74 new_class.__repr__ = sn_repr 75 76 return new_class 77# pylint: enable=docstring-misnamed-args 78 79 80def GroupByKey(input_iter, key): 81 """Split an iterable of dicts, based on value of a key. 82 83 GroupByKey([{'a': 1}, {'a': 2}, {'a': 1, 'b': 2}], 'a') => 84 {1: [{'a': 1}, {'a': 1, 'b': 2}], 2: [{'a': 2}]} 85 86 Args: 87 input_iter: An iterable of dicts. 88 key: A string specifying the key name to split by. 89 90 Returns: 91 A dictionary, mapping from each unique value for |key| that 92 was encountered in |input_iter| to a list of entries that had 93 that value. 94 """ 95 split_dict = dict() 96 for entry in input_iter: 97 split_dict.setdefault(entry.get(key), []).append(entry) 98 return split_dict 99 100 101def GroupNamedtuplesByKey(input_iter, key): 102 """Split an iterable of namedtuples, based on value of a key. 103 104 Args: 105 input_iter: An iterable of namedtuples. 106 key: A string specifying the key name to split by. 107 108 Returns: 109 A dictionary, mapping from each unique value for |key| that 110 was encountered in |input_iter| to a list of entries that had 111 that value. 112 """ 113 split_dict = {} 114 for entry in input_iter: 115 split_dict.setdefault(getattr(entry, key, None), []).append(entry) 116 return split_dict 117 118 119def InvertDictionary(origin_dict): 120 """Invert the key value mapping in the origin_dict. 121 122 Given an origin_dict {'key1': {'val1', 'val2'}, 'key2': {'val1', 'val3'}, 123 'key3': {'val3'}}, the returned inverted dict will be 124 {'val1': {'key1', 'key2'}, 'val2': {'key1'}, 'val3': {'key2', 'key3'}} 125 126 Args: 127 origin_dict: A dict mapping each key to a group (collection) of values. 128 129 Returns: 130 An inverted dict mapping each key to a set of its values. 131 """ 132 new_dict = {} 133 for origin_key, origin_values in origin_dict.items(): 134 for origin_value in origin_values: 135 new_dict.setdefault(origin_value, set()).add(origin_key) 136 137 return new_dict 138