Merge remote-tracking branch 'scottkmaxwell/master' into item_sort_key

4 files changed, 77 insertions, 16 deletions

diff --git a/simplejson/__init__.py b/simplejson/__init__.py
index 4a5e5df..3de54de 100644
--- a/simplejson/__init__.py
+++ b/simplejson/__init__.py
@@ -101,7 +101,7 @@ __version__ = '2.4.0'
 __all__ = [
     'dump', 'dumps', 'load', 'loads',
     'JSONDecoder', 'JSONDecodeError', 'JSONEncoder',
-    'OrderedDict',
+    'OrderedDict', 'simple_first'
 ]
 
 __author__ = 'Bob Ippolito <bob@redivi.com>'
@@ -139,13 +139,14 @@ _default_encoder = JSONEncoder(
     namedtuple_as_object=True,
     tuple_as_array=True,
     bigint_as_string=False,
+    item_sort_key=None,
 )
 
 def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
         allow_nan=True, cls=None, indent=None, separators=None,
         encoding='utf-8', default=None, use_decimal=True,
         namedtuple_as_object=True, tuple_as_array=True,
-        bigint_as_string=False, **kw):
+        bigint_as_string=False, item_sort_key=None, **kw):
     """Serialize ``obj`` as a JSON formatted stream to ``fp`` (a
     ``.write()``-supporting file-like object).
 
@@ -200,6 +201,10 @@ def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
     lossy operation that will not round-trip correctly and should be used
     sparingly.
 
+    If specified, *item_sort_key* is a callable used to sort the items in
+    each dictionary. This is useful if you want to sort items other than
+    in alphabetical order by key.
+
     To use a custom ``JSONEncoder`` subclass (e.g. one that overrides the
     ``.default()`` method to serialize additional types), specify it with
     the ``cls`` kwarg.
@@ -211,7 +216,7 @@ def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
         cls is None and indent is None and separators is None and
         encoding == 'utf-8' and default is None and use_decimal
         and namedtuple_as_object and tuple_as_array
-        and not bigint_as_string and not kw):
+        and not bigint_as_string and not item_sort_key and not kw):
         iterable = _default_encoder.iterencode(obj)
     else:
         if cls is None:
@@ -223,6 +228,7 @@ def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
             namedtuple_as_object=namedtuple_as_object,
             tuple_as_array=tuple_as_array,
             bigint_as_string=bigint_as_string,
+            item_sort_key=item_sort_key,
             **kw).iterencode(obj)
     # could accelerate with writelines in some versions of Python, at
     # a debuggability cost
@@ -233,8 +239,8 @@ def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
 def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
         allow_nan=True, cls=None, indent=None, separators=None,
         encoding='utf-8', default=None, use_decimal=True,
-        namedtuple_as_object=True,
-        tuple_as_array=True, bigint_as_string=False,
+        namedtuple_as_object=True, tuple_as_array=True,
+        bigint_as_string=False, item_sort_key=None,
         **kw):
     """Serialize ``obj`` to a JSON formatted ``str``.
 
@@ -285,6 +291,10 @@ def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
     or lower than -2**53 will be encoded as strings. This is to avoid the
     rounding that happens in Javascript otherwise.
 
+    If specified, *item_sort_key* is a callable used to sort the items in
+    each dictionary. This is useful if you want to sort items other than
+    in alphabetical order by key.
+
     To use a custom ``JSONEncoder`` subclass (e.g. one that overrides the
     ``.default()`` method to serialize additional types), specify it with
     the ``cls`` kwarg.
@@ -296,7 +306,7 @@ def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
         cls is None and indent is None and separators is None and
         encoding == 'utf-8' and default is None and use_decimal
         and namedtuple_as_object and tuple_as_array
-        and not bigint_as_string and not kw):
+        and not bigint_as_string and not item_sort_key and not kw):
         return _default_encoder.encode(obj)
     if cls is None:
         cls = JSONEncoder
@@ -308,6 +318,7 @@ def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
         namedtuple_as_object=namedtuple_as_object,
         tuple_as_array=tuple_as_array,
         bigint_as_string=bigint_as_string,
+        item_sort_key=item_sort_key,
         **kw).encode(obj)
 
 
@@ -479,3 +490,9 @@ def _toggle_speedups(enabled):
        encoding='utf-8',
        default=None,
    )
+
+def simple_first(kv):
+    """Helper function to pass to item_sort_key to sort simple
+    elements to the top, then container elements.
+    """
+    return isinstance(kv[1],(list,dict,tuple)), kv[0]
diff --git a/simplejson/_speedups.c b/simplejson/_speedups.c
index a0e5667..f6eef80 100644
--- a/simplejson/_speedups.c
+++ b/simplejson/_speedups.c
@@ -90,6 +90,7 @@ typedef struct _PyEncoderObject {
     int namedtuple_as_object;
     int tuple_as_array;
     int bigint_as_string;
+    PyObject *item_sort_key;
 } PyEncoderObject;
 
 static PyMemberDef encoder_members[] = {
@@ -102,6 +103,7 @@ static PyMemberDef encoder_members[] = {
     {"sort_keys", T_OBJECT, offsetof(PyEncoderObject, sort_keys), READONLY, "sort_keys"},
     {"skipkeys", T_OBJECT, offsetof(PyEncoderObject, skipkeys), READONLY, "skipkeys"},
     {"key_memo", T_OBJECT, offsetof(PyEncoderObject, key_memo), READONLY, "key_memo"},
+    {"item_sort_key", T_OBJECT, offsetof(PyEncoderObject, item_sort_key), READONLY, "item_sort_key"},
     {NULL}
 };
 
@@ -2050,6 +2052,7 @@ encoder_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         s->sort_keys = NULL;
         s->skipkeys = NULL;
         s->key_memo = NULL;
+        s->item_sort_key = NULL;
     }
     return (PyObject *)s;
 }
@@ -2058,19 +2061,19 @@ static int
 encoder_init(PyObject *self, PyObject *args, PyObject *kwds)
 {
     /* initialize Encoder object */
-    static char *kwlist[] = {"markers", "default", "encoder", "indent", "key_separator", "item_separator", "sort_keys", "skipkeys", "allow_nan", "key_memo", "use_decimal", "namedtuple_as_object", "tuple_as_array", "bigint_as_string", NULL};
+    static char *kwlist[] = {"markers", "default", "encoder", "indent", "key_separator", "item_separator", "sort_keys", "skipkeys", "allow_nan", "key_memo", "use_decimal", "namedtuple_as_object", "tuple_as_array", "bigint_as_string", "item_sort_key", NULL};
 
     PyEncoderObject *s;
     PyObject *markers, *defaultfn, *encoder, *indent, *key_separator;
-    PyObject *item_separator, *sort_keys, *skipkeys, *allow_nan, *key_memo, *use_decimal, *namedtuple_as_object, *tuple_as_array, *bigint_as_string;
+    PyObject *item_separator, *sort_keys, *skipkeys, *allow_nan, *key_memo, *use_decimal, *namedtuple_as_object, *tuple_as_array, *bigint_as_string, *item_sort_key;
 
     assert(PyEncoder_Check(self));
     s = (PyEncoderObject *)self;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OOOOOOOOOOOOOO:make_encoder", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OOOOOOOOOOOOOOO:make_encoder", kwlist,
         &markers, &defaultfn, &encoder, &indent, &key_separator, &item_separator,
         &sort_keys, &skipkeys, &allow_nan, &key_memo, &use_decimal,
-        &namedtuple_as_object, &tuple_as_array, &bigint_as_string))
+        &namedtuple_as_object, &tuple_as_array, &bigint_as_string, &item_sort_key))
         return -1;
 
     s->markers = markers;
@@ -2088,6 +2091,7 @@ encoder_init(PyObject *self, PyObject *args, PyObject *kwds)
     s->namedtuple_as_object = PyObject_IsTrue(namedtuple_as_object);
     s->tuple_as_array = PyObject_IsTrue(tuple_as_array);
     s->bigint_as_string = PyObject_IsTrue(bigint_as_string);
+    s->item_sort_key = item_sort_key;
 
     Py_INCREF(s->markers);
     Py_INCREF(s->defaultfn);
@@ -2098,6 +2102,7 @@ encoder_init(PyObject *self, PyObject *args, PyObject *kwds)
     Py_INCREF(s->sort_keys);
     Py_INCREF(s->skipkeys);
     Py_INCREF(s->key_memo);
+    Py_INCREF(s->item_sort_key);
     return 0;
 }
 
@@ -2356,7 +2361,14 @@ encoder_listencode_dict(PyEncoderObject *s, PyObject *rval, PyObject *dct, Py_ss
         */
     }
 
-    if (PyObject_IsTrue(s->sort_keys)) {
+	if (PyCallable_Check(s->item_sort_key)) {
+        if (PyDict_CheckExact(dct))
+            items = PyDict_Items(dct);
+        else
+            items = PyMapping_Items(dct);
+        PyObject_CallMethod(items, "sort", "OO", Py_None, s->item_sort_key);
+	}
+    else if (PyObject_IsTrue(s->sort_keys)) {
         /* First sort the keys then replace them with (key, value) tuples. */
         Py_ssize_t i, nitems;
         if (PyDict_CheckExact(dct))
@@ -2616,6 +2628,7 @@ encoder_traverse(PyObject *self, visitproc visit, void *arg)
     Py_VISIT(s->sort_keys);
     Py_VISIT(s->skipkeys);
     Py_VISIT(s->key_memo);
+    Py_VISIT(s->item_sort_key);
     return 0;
 }
 
@@ -2635,6 +2648,7 @@ encoder_clear(PyObject *self)
     Py_CLEAR(s->sort_keys);
     Py_CLEAR(s->skipkeys);
     Py_CLEAR(s->key_memo);
+    Py_CLEAR(s->item_sort_key);
     return 0;
 }
 
diff --git a/simplejson/encoder.py b/simplejson/encoder.py
index d13d26a..1790e70 100644
--- a/simplejson/encoder.py
+++ b/simplejson/encoder.py
@@ -107,7 +107,8 @@ class JSONEncoder(object):
             check_circular=True, allow_nan=True, sort_keys=False,
             indent=None, separators=None, encoding='utf-8', default=None,
             use_decimal=True, namedtuple_as_object=True,
-            tuple_as_array=True, bigint_as_string=False):
+            tuple_as_array=True, bigint_as_string=False,
+            item_sort_key=None):
         """Constructor for JSONEncoder, with sensible defaults.
 
         If skipkeys is false, then it is a TypeError to attempt
@@ -164,6 +165,10 @@ class JSONEncoder(object):
         If bigint_as_string is true (not the default), ints 2**53 and higher
         or lower than -2**53 will be encoded as strings. This is to avoid the
         rounding that happens in Javascript otherwise.
+
+        If specified, item_sort_key is a callable used to sort the items in
+        each dictionary. This is useful if you want to sort items other than
+        in alphabetical order by key.
         """
 
         self.skipkeys = skipkeys
@@ -175,6 +180,7 @@ class JSONEncoder(object):
         self.namedtuple_as_object = namedtuple_as_object
         self.tuple_as_array = tuple_as_array
         self.bigint_as_string = bigint_as_string
+        self.item_sort_key = item_sort_key
         if indent is not None and not isinstance(indent, basestring):
             indent = indent * ' '
         self.indent = indent
@@ -291,14 +297,14 @@ class JSONEncoder(object):
                 self.key_separator, self.item_separator, self.sort_keys,
                 self.skipkeys, self.allow_nan, key_memo, self.use_decimal,
                 self.namedtuple_as_object, self.tuple_as_array,
-                self.bigint_as_string)
+                self.bigint_as_string, self.item_sort_key)
         else:
             _iterencode = _make_iterencode(
                 markers, self.default, _encoder, self.indent, floatstr,
                 self.key_separator, self.item_separator, self.sort_keys,
                 self.skipkeys, _one_shot, self.use_decimal,
                 self.namedtuple_as_object, self.tuple_as_array,
-                self.bigint_as_string)
+                self.bigint_as_string, self.item_sort_key)
         try:
             return _iterencode(o, 0)
         finally:
@@ -335,7 +341,7 @@ class JSONEncoderForHTML(JSONEncoder):
 def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
         _key_separator, _item_separator, _sort_keys, _skipkeys, _one_shot,
         _use_decimal, _namedtuple_as_object, _tuple_as_array,
-        _bigint_as_string,
+        _bigint_as_string, _item_sort_key,
         ## HACK: hand-optimized bytecode; turn globals into locals
         False=False,
         True=True,
@@ -437,7 +443,10 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
             newline_indent = None
             item_separator = _item_separator
         first = True
-        if _sort_keys:
+        if _item_sort_key:
+            items = dct.items()
+            items.sort(key=_item_sort_key)
+        elif _sort_keys:
             items = dct.items()
             items.sort(key=lambda kv: kv[0])
         else:
diff --git a/simplejson/tests/test_item_sort_key.py b/simplejson/tests/test_item_sort_key.py
new file mode 100644
index 0000000..f630d8e
--- /dev/null
+++ b/simplejson/tests/test_item_sort_key.py
@@ -0,0 +1,21 @@
+from unittest import TestCase
+
+import simplejson as json
+from operator import itemgetter
+
+class TestItemSortKey(TestCase):
+    def test_simple_first(self):
+        a={'a': 1, 'c': 5, 'jack': 'jill', 'pick': 'axe', 'array': [1, 5, 6, 9], 'tuple': (83, 12, 3), 'crate': 'dog', 'zeak': 'oh'}
+        self.assertEquals(
+            '{"a": 1, "c": 5, "crate": "dog", "jack": "jill", "pick": "axe", "zeak": "oh", "array": [1, 5, 6, 9], "tuple": [83, 12, 3]}',
+            json.dumps(a,item_sort_key=json.simple_first))
+
+    def test_case(self):
+        a={'a': 1, 'c': 5, 'Jack': 'jill', 'pick': 'axe', 'Array': [1, 5, 6, 9], 'tuple': (83, 12, 3), 'crate': 'dog', 'zeak': 'oh'}
+        self.assertEquals(
+            '{"Array": [1, 5, 6, 9], "Jack": "jill", "a": 1, "c": 5, "crate": "dog", "pick": "axe", "tuple": [83, 12, 3], "zeak": "oh"}',
+            json.dumps(a,item_sort_key=itemgetter(0)))
+        self.assertEquals(
+            '{"a": 1, "Array": [1, 5, 6, 9], "c": 5, "crate": "dog", "Jack": "jill", "pick": "axe", "tuple": [83, 12, 3], "zeak": "oh"}',
+            json.dumps(a,item_sort_key=lambda kv: kv[0].lower()))
+