-

cool-RR · cool-RR · commit 3f354473a135 · 2014-01-19T22:34:35.000+02:00
diff --git a/source_py2/python_toolbox/collection_tools.py b/source_py2/python_toolbox/collection_tools.py
@@ -0,0 +1,24 @@
+# Copyright 2009-2014 Ram Rachum.
+# This program is distributed under the MIT license.
+
+'''Defines various tools related to data structures.'''
+
+import collections
+import itertools
+import numbers
+
+from python_toolbox import nifty_collections
+
+
+@nifty_collections.LazyTuple.factory
+def get_all_contained_counters(counter):
+    assert isinstance(counter, collections.Counter)
+    counter_type = type(counter)
+    keys, amounts = zip(
+        *((key, amount) for key, amount in counter.items() if amount)
+    )
+    assert all(isinstance(amount, numbers.Integral) for amount in amounts)
+    amounts_tuples = \
+               itertools.product(*map(lambda amount: xrange(amount+1), amounts))
+    for amounts_tuple in amounts_tuples:
+        yield counter_type(dict(zip(keys, amounts_tuple)))
diff --git a/source_py2/test_python_toolbox/test_collection_tools/__init__.py b/source_py2/test_python_toolbox/test_collection_tools/__init__.py
diff --git a/source_py2/test_python_toolbox/test_collection_tools/test_get_all_contained_counters.py b/source_py2/test_python_toolbox/test_collection_tools/test_get_all_contained_counters.py
@@ -0,0 +1,37 @@
+# Copyright 2009-2014 Ram Rachum.
+# This program is distributed under the MIT license.
+
+import collections
+
+from python_toolbox import nifty_collections
+
+from python_toolbox import collection_tools
+
+
+def test():
+    
+    counter = collections.Counter('abracadabra')
+    assert counter == collections.Counter({'a': 5, 'b': 2, 'r': 2, 'c': 1,
+                                           'd': 1,})
+    all_contained_counters = \
+                           collection_tools.get_all_contained_counters(counter)
+    
+    assert isinstance(all_contained_counters, nifty_collections.LazyTuple)
+    
+    assert not all_contained_counters.collected_data
+    
+    # Now we'll exhaust `all_contained_counters`:
+    assert len(all_contained_counters) == 6 * 3 * 3 * 2 * 2
+    assert all_contained_counters.exhausted
+    
+    assert counter in all_contained_counters
+    assert collections.Counter({'a': 0, 'b': 0, 'r': 0, 'c': 0, 'd': 0,}) in \
+                                                         all_contained_counters
+    assert collections.Counter({'a': 1, 'b': 1, 'r': 1, 'c': 1, 'd': 1,}) in \
+                                                         all_contained_counters
+    assert collections.Counter({'a': 2, 'b': 2, 'r': 2, 'c': 1, 'd': 1,}) in \
+                                                         all_contained_counters
+    assert collections.Counter({'a': 4, 'b': 2, 'r': 2, 'c': 1, 'd': 1,}) in \
+                                                         all_contained_counters
+    assert all(isinstance(item, collections.Counter) for item in
+                                                        all_contained_counters)
