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