Merge pull request #1 from jsiirola/sorted-robust-write

Extend sorted_robust API

pyomo/core/base/block.py

@@ -1359,7 +1359,7 @@ def _component_data_iter(self, ctype=None, active=None, sort=False):
                 _items = ((None, comp),)
 
             if _sort_indices:
-                _items = sorted_robust(_items)
+                _items = sorted_robust(_items, key=itemgetter(0))
             if active is None or not isinstance(comp, ActiveIndexedComponent):
                 for idx, compData in _items:
                     yield (name, idx), compData

pyomo/core/base/misc.py

@@ -14,6 +14,8 @@
 import sys
 import types
 
+from pyomo.core.expr import native_numeric_types
+
 logger = logging.getLogger('pyomo.core')
 
 
@@ -110,8 +112,11 @@ class _robust_sort_keyfcn(object):
     _typemap without resorting to global variables.
 
     """
-    def __init__(self):
-        self._typemap = {tuple: (3, tuple.__name__)}
+    def __init__(self, key=None):
+        # sort all native numeric types as if they were floats
+        self._typemap = {t: (1, float.__name__) for t in native_numeric_types}
+        self._typemap[tuple] =(3, tuple.__name__)
+        self._key = key
 
     def __call__(self, val):
         """Generate a tuple ( str(type_name), val ) for sorting the value.
@@ -122,6 +127,9 @@ def __call__(self, val):
         argument of the sort key.
 
         """
+        if self._key is not None:
+            val = self._key(val)
+
         try:
             i, _typename = self._typemap[val.__class__]
         except KeyError:
@@ -155,7 +163,7 @@ def __call__(self, val):
             return _typename, id(val)
 
 
-def sorted_robust(arg):
+def sorted_robust(arg, key=None, reverse=False):
     """Utility to sort an arbitrary iterable.
 
     This returns the sorted(arg) in a consistent order by first tring
@@ -166,16 +174,16 @@ def sorted_robust(arg):
     """
     # It is possible that arg is a generator.  We need to cache the
     # elements returned by the generator in case 'sorted' raises an
-    # exception (this ensures we don't loose any elements).  Howevver,
+    # exception (this ensures we don't lose any elements).  However,
     # if we were passed a list, we do not want to make an unnecessary
     # copy.  Tuples are OK because tuple(a) will not copy a if it is
     # already a tuple.
     if type(arg) is not list:
         arg = tuple(arg)
     try:
-        return sorted(arg)
+        return sorted(arg, key=key, reverse=reverse)
     except:
-        return sorted(arg, key=_robust_sort_keyfcn())
+        return sorted(arg, key=_robust_sort_keyfcn(key), reverse=reverse)
 
 
 def _to_ustr(obj):

pyomo/core/tests/unit/test_base_misc.py

@@ -38,6 +38,10 @@ def test_sorted_robust(self):
         a = sorted_robust([3, 2, 1])
         self.assertEqual(a, [1, 2, 3])
 
+        # Testthat ints and floats are sorted as "numbers"
+        a = sorted_robust([3, 2.1, 1])
+        self.assertEqual(a, [1, 2.1, 3])
+
         a = sorted_robust([3, '2', 1])
         self.assertEqual(a, [1, 3, '2'])
 

pyomo/core/tests/unit/test_block.py

@@ -342,6 +342,15 @@ def test_Block(self):
         self.assertEqual(sorted([id(comp) for comp in model.block_data_objects(sort=False)]),
                          sorted([id(comp) for comp in [model,]+model.component_data_lists[Block]]))
 
+    def test_mixed_index_type(self):
+        m = ConcreteModel()
+        m.I = Set(initialize=[1,'1',3.5,4])
+        m.x = Var(m.I)
+        v = list(m.component_data_objects(Var, sort=True))
+        self.assertEqual(len(v), 4)
+        for a,b in zip([m.x[1], m.x[3.5], m.x[4], m.x['1']], v):
+            self.assertIs(a, b)
+
 
 class HierarchicalModel(object):
     def __init__(self):