1""" 2Various useful utilities, mostly taken from the ASPN Python cookbook. 3""" 4 5import types 6seq_types = type(()), type([]) 7 8 9def flatten(*args): 10 for arg in args: 11 if type(arg) in seq_types: 12 for elem in arg: 13 yield from flatten(elem) 14 else: 15 yield arg 16 17 18def cross_lists(*sets): 19 """Return the cross product of the arguments""" 20 wheels = [iter(_) for _ in sets] 21 digits = [next(it) for it in wheels] 22 while True: 23 yield digits[:] 24 for i in range(len(digits)-1, -1, -1): 25 try: 26 digits[i] = next(wheels[i]) 27 break 28 except StopIteration: 29 wheels[i] = iter(sets[i]) 30 digits[i] = next(wheels[i]) 31 else: 32 break 33 34# Cached / memoized methods 35 36 37def cachedmethod(function): 38 return types.MethodType(Memoize(function), None) 39 40 41class Memoize: 42 def __init__(self, function): 43 self._cache = {} 44 self._callable = function 45 46 def __call__(self, *args, **kwds): 47 cache = self._cache 48 key = self._getKey(*args, **kwds) 49 try: 50 return cache[key] 51 except KeyError: 52 cachedValue = cache[key] = self._callable(*args, **kwds) 53 return cachedValue 54 55 def _getKey(self, *args, **kwds): 56 return kwds and (args, ImmutableDict(kwds)) or args 57 58 59class memoized: 60 """Decorator that caches a function's return value each time it is called. 61 If called later with the same arguments, the cached value is returned, and 62 not re-evaluated. 63 """ 64 65 def __init__(self, func): 66 self.func = func 67 self.cache = {} 68 69 def __call__(self, *args): 70 try: 71 return self.cache[args] 72 except KeyError: 73 self.cache[args] = value = self.func(*args) 74 return value 75 except TypeError: 76 # uncachable -- for instance, passing a list as an argument. 77 # Better to not cache than to blow up entirely. 78 return self.func(*args) 79 80 def __repr__(self): 81 """Return the function's docstring.""" 82 return self.func.__doc__ 83 84 85class ImmutableDict(dict): 86 '''A hashable dict.''' 87 88 def __init__(self, *args, **kwds): 89 dict.__init__(self, *args, **kwds) 90 91 def __setitem__(self, key, value): 92 raise NotImplementedError("dict is immutable") 93 94 def __delitem__(self, key): 95 raise NotImplementedError("dict is immutable") 96 97 def clear(self): 98 raise NotImplementedError("dict is immutable") 99 100 def setdefault(self, k, default=None): 101 raise NotImplementedError("dict is immutable") 102 103 def popitem(self): 104 raise NotImplementedError("dict is immutable") 105 106 def update(self, other): 107 raise NotImplementedError("dict is immutable") 108 109 def __hash__(self): 110 return hash(tuple(self.items())) 111