Friendsnippets

 1 """A function to enumerate the items of a set product.
 2 """
 3 __docformat__ = 'reStructuredText'
 4 
 5 
 6 def enumerate_set_product(*args):
 7     """Iterate over all elements in the cartesian product of its arguments.
 8     Each argument to this variadic function *must* be a sequence.
 9     
10     Examples::
11       >>> list(enumerate_set_product())
12       [[]]
13       >>> list(enumerate_set_product([1]))
14       [[1]]
15       >>> list(enumerate_set_product([1],[1]))
16       [[1, 1]]
17       >>> list(enumerate_set_product([1,2],[1]))
18       [[1, 1], [2, 1]]
19       >>> list(enumerate_set_product([1,2],[1,2]))
20       [[1, 1], [2, 1], [1, 2], [2, 2]]
21     """
22 
23     # `assert` guard against invocation with wrong arguments.
24     def __all_items_are_sequences(s):
25         from operator import and_, isSequenceType
26         return reduce(and_,
27                       [ isSequenceType(item) for item in s],
28                       True)
29     assert __all_items_are_sequences(args), \
30            "All arguments to `enumerate_set_products` must be sequences."
31     
32     if len(args) == 0:
33         yield []
34     else:
35         L = len(args)
36         M = [ len(s)-1 for s in args ]
37         m = [0] * L
38         i = 0
39         while i < L:
40             # return element corresponding to current multi-index
41             yield [ s[m[i]] for (i,s) in enumerate(args) ]
42             # advance multi-index
43             i = 0
44             while (i < L):
45                 if m[i] == M[i]:
46                     m[i] = 0
47                     i += 1
48                 else:
49                     m[i] += 1
50                     break
51 
52 
53 def recursively_enumerate_set_product(*args):
54     """Iterate over all elements in the cartesian product of its arguments.
55      Each argument to this variadic function *must* be a sequence.
56 
57      You would probably prefer `enumerate_set_product`, which is faster.
58 
59     Examples::
60       >>> list(recursively_enumerate_set_product())
61       [[]]
62       >>> list(recursively_enumerate_set_product([1]))
63       [[1]]
64       >>> list(recursively_enumerate_set_product([1],[1]))
65       [[1, 1]]
66       >>> list(recursively_enumerate_set_product([1,2],[1]))
67       [[1, 1], [2, 1]]
68       >>> list(recursively_enumerate_set_product([1,2],[1,2]))
69       [[1, 1], [2, 1], [1, 2], [2, 2]]
70     """
71 
72     # `assert` guard against invocation with wrong arguments.
73     def __all_items_are_sequences(s):
74         from operator import and_, isSequenceType
75         return reduce(and_,
76                       [ isSequenceType(item) for item in s],
77                       True)
78     assert __all_items_are_sequences(args), \
79            "All arguments to `enumerate_set_products` must be sequences."
80     
81     if len(args) == 0:
82         yield []
83     else:
84         for i in args[-1]:
85             for js in recursively_enumerate_set_product(*args[:-1]):
86                 yield js+[i]
87 
88 
89 ## main: run tests
90 
91 if "__main__" == __name__:
92     import doctest
93     doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)