-

cool-RR · cool-RR · commit 2b8d34689670 · 2012-12-23T18:49:39.000+02:00
diff --git a/python_toolbox/dict_tools.py b/python_toolbox/dict_tools.py
@@ -3,6 +3,8 @@
 
 '''Defines several functions that may be useful when working with dicts.'''
 
+from python_toolbox import cute_iter_tools
+
 
 def filter_items(d, condition, force_dict_type=None):
     '''
@@ -67,9 +69,31 @@ def reverse_with_set_values(d):
         reverse_with_set_values({1: 2, 3: 4, 'meow': 2}) == \
             {2: set([1, 'meow']), 4: set([3])}
             
+    Instead of a dict you may also input a tuple in which the first item is an
+    iterable and the second item is either a key function or an attribute name.
+    A dict will be constructed from these and used.
     '''
+    ### Pre-processing input: #################################################
+    #                                                                         #
+    if hasattr(d, 'items'): # `d` is a dict
+        fixed_dict = d
+    else: # `d` is not a dict
+        assert cute_iter_tools.is_iterable(d) and len(d) == 2
+        iterable, key_function_or_attribute_name = d
+        assert cute_iter_tools.is_iterable(iterable)
+        if callable(key_function_or_attribute_name):
+            key_function = key_function_or_attribute_name
+        else:
+            assert isinstance(key_function_or_attribute_name, basestring)
+            key_function = \
+                       lambda key: getattr(key, key_function_or_attribute_name)
+            
+        fixed_dict = dict((key, key_function(key)) for key in iterable)
+    #                                                                         #
+    ### Finished pre-processing input. ########################################
+    
     new_dict = {}
-    for key, value in d.iteritems():
+    for key, value in fixed_dict.iteritems():
         if value not in new_dict:
             new_dict[value] = []
         new_dict[value].append(key)
diff --git a/test_python_toolbox/test_dict_tools/test_reverse_with_set_values.py b/test_python_toolbox/test_dict_tools/test_reverse_with_set_values.py
@@ -9,4 +9,12 @@
 def test():
     '''Test the basic workings of `reverse_with_set_values`.'''
     assert dict_tools.reverse_with_set_values({1: 2, 3: 4, 'meow': 2}) == \
-                                             {2: set([1, 'meow']), 4: set([3])}
+                                             {2: set([1, 'meow']), 4: set([3])}
+    
+def test_iterable_input():
+    assert dict_tools.reverse_with_set_values((range(1, 5), str)) == \
+        {'1': set([1]), '2': set([2]), '3': set([3]), '4': set([4]),}
+    
+    assert dict_tools.reverse_with_set_values(([1, 2+3j, 4, 5-6j], 'imag')) \
+                           == {0: set([1, 4]), 3: set([2+3j]), -6: set([5-6j])}
+    
