-

cool-RR · cool-RR · commit 886f45c2e576 · 2012-12-24T01:41:10.000+02:00
diff --git a/python_toolbox/comparison_tools.py b/python_toolbox/comparison_tools.py
@@ -54,3 +54,15 @@ def total_ordering(cls):
             setattr(cls, opname, opfunc)
     return cls
 
+
+def process_key_function_or_attribute_name(key_function_or_attribute_name):
+    if key_function_or_attribute_name is None:
+        return None
+    elif callable(key_function_or_attribute_name):
+        return key_function_or_attribute_name
+    else:
+        assert isinstance(key_function_or_attribute_name, basestring)
+        return lambda key: getattr(key, key_function_or_attribute_name)
+
+
+
diff --git a/python_toolbox/dict_tools.py b/python_toolbox/dict_tools.py
@@ -4,6 +4,7 @@
 '''Defines several functions that may be useful when working with dicts.'''
 
 from python_toolbox import cute_iter_tools
+from python_toolbox import comparison_tools
 
 
 def filter_items(d, condition, force_dict_type=None):
@@ -60,7 +61,7 @@ def show(thing, indent=0):
     return temp2
     
 
-def reverse_with_set_values(d):
+def reverse_with_set_values(d, sort=False):
     '''
     Reverse the dict, with the values of the new dict being sets.
     
@@ -72,6 +73,10 @@ def reverse_with_set_values(d):
     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.
+    
+    If you'd like the result dict to be sorted, pass `sort=True`, and you'll
+    get a sorted `OrderedDict`. You can also specify the sorting key function
+    or attribute name as the `sort` argument.
     '''
     ### Pre-processing input: #################################################
     #                                                                         #
@@ -81,13 +86,9 @@ def reverse_with_set_values(d):
         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)
-            
+        key_function = comparison_tools.process_key_function_or_attribute_name(
+            key_function_or_attribute_name
+        )
         fixed_dict = dict((key, key_function(key)) for key in iterable)
     #                                                                         #
     ### Finished pre-processing input. ########################################
@@ -102,7 +103,19 @@ def reverse_with_set_values(d):
     for key, value in new_dict.copy().iteritems():
         new_dict[key] = set(value)
         
-    return new_dict
+    if sort:
+        from python_toolbox import nifty_collections
+        ordered_dict = nifty_collections.OrderedDict(new_dict)
+        if callable(sort) or isinstance(sort, basestring):
+            key_function = comparison_tools. \
+                                   process_key_function_or_attribute_name(sort)
+        else:
+            assert sort is True
+            key_function = None
+        ordered_dict.sort(key_function)
+        return ordered_dict
+    else:
+        return new_dict
 
 
 def devour_items(d):
diff --git a/python_toolbox/nifty_collections/ordered_dict.py b/python_toolbox/nifty_collections/ordered_dict.py
@@ -7,9 +7,10 @@
 See its documentation for more information.
 '''
 
-
 from UserDict import DictMixin
 
+from python_toolbox import comparison_tools
+
 
 class OrderedDict(dict, DictMixin):
     '''Dict that maintains order of items.'''
@@ -168,9 +169,11 @@ def sort(self, key=None):
         The optional `key` argument, (not to be confused with the dictionary
         keys,) will be passed to the `sorted` function as a key function.
         '''
-        sorted_keys = sorted(self.keys(), key=key)
-        for key in sorted_keys[1:]:
-            self.move_to_end(key)
+        key_function = \
+                   comparison_tools.process_key_function_or_attribute_name(key)
+        sorted_keys = sorted(self.keys(), key=key_function)
+        for key_ in sorted_keys[1:]:
+            self.move_to_end(key_)
         
             
     def index(self, 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
@@ -4,6 +4,7 @@
 '''Test package for `dict_tools.reverse_with_set_values`.'''
 
 from python_toolbox import dict_tools
+from python_toolbox import nifty_collections
 
 
 def test():
@@ -17,4 +18,14 @@ def test_iterable_input():
     
     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])}
+    
+    
+def test_ordered_dict_output():
+    assert dict_tools.reverse_with_set_values({1: 2, 3: 4, 'meow': 2},
+                                              sort=True) == \
+    nifty_collections.OrderedDict([(2, set([1, 'meow'])), (4, set([3]))])
+    
+    assert dict_tools.reverse_with_set_values({1: 2, 3: 4, 'meow': 2},
+                                              sort=lambda x: -x) == \
+    nifty_collections.OrderedDict([(4, set([3])), (2, set([1, 'meow']))])
     
diff --git a/test_python_toolbox/test_nifty_collections/test_ordered_dict/test.py b/test_python_toolbox/test_nifty_collections/test_ordered_dict/test.py
@@ -37,6 +37,13 @@ def test_sort():
     assert ordered_dict == ordered_dict_copy
     
     
+    second_ordered_dict = OrderedDict(((1+2j, 'b'), (2+3j, 'c'), (3+1j, 'a')))
+    second_ordered_dict.sort('imag')
+    assert second_ordered_dict == \
+                           OrderedDict(((3+1j, 'a'), (1+2j, 'b'), (2+3j, 'c')))
+
+    
+    
 def test_index():
     '''Test the `OrderedDict.index` method.'''
     ordered_dict = OrderedDict(((1, 'a'), (2, 'b'), (3, 'c')))
