Implement for_json kwarg; pure Python version; add tests; update documentation.

5 files changed, 210 insertions, 39 deletions

diff --git a/index.rst b/index.rst
index 4dd8a58..c9dd246 100644
--- a/index.rst
+++ b/index.rst
@@ -245,6 +245,13 @@ Basic Usage
       strings, to avoid comparison of heterogeneously typed objects
       (since this does not work in Python 3.3+)
 
+   If *for_json* is true (not the default), objects with a ``for_json()``
+   method will use the return value of that method for encoding as JSON instead
+   of the object.
+
+   .. versionchanged:: 3.2.0
+      *for_json* is new in 3.2.0.
+
     .. note::
 
         JSON is not a framed protocol so unlike :mod:`pickle` or :mod:`marshal` it
@@ -252,7 +259,7 @@ Basic Usage
         container protocol to delimit them.
 
 
-.. function:: dumps(obj[, skipkeys[, ensure_ascii[, check_circular[, allow_nan[, cls[, indent[, separators[, encoding[, default[, use_decimal[, namedtuple_as_object[, tuple_as_array[, bigint_as_string[, sort_keys[, item_sort_key[, **kw]]]]]]]]]]]]]]]])
+.. function:: dumps(obj[, skipkeys[, ensure_ascii[, check_circular[, allow_nan[, cls[, indent[, separators[, encoding[, default[, use_decimal[, namedtuple_as_object[, tuple_as_array[, bigint_as_string[, sort_keys[, item_sort_key[, for_json[, **kw]]]]]]]]]]]]]]]]])
 
    Serialize *obj* to a JSON formatted :class:`str`.
 
@@ -456,7 +463,7 @@ Encoders and decoders
       :exc:`JSONDecodeError` will be raised if the given JSON
       document is not valid.
 
-.. class:: JSONEncoder([skipkeys[, ensure_ascii[, check_circular[, allow_nan[, sort_keys[, indent[, separators[, encoding[, default[, use_decimal[, namedtuple_as_object[, tuple_as_array[, bigint_as_string[, item_sort_key]]]]]]]]]]]]])
+.. class:: JSONEncoder([skipkeys[, ensure_ascii[, check_circular[, allow_nan[, sort_keys[, indent[, separators[, encoding[, default[, use_decimal[, namedtuple_as_object[, tuple_as_array[, bigint_as_string[, item_sort_key[, for_json]]]]]]]]]]]]]])
 
    Extensible JSON encoder for Python data structures.
 
@@ -587,6 +594,12 @@ Encoders and decoders
    .. versionchanged:: 2.4.0
      *bigint_as_string* is new in 2.4.0.
 
+   If *for_json* is true (default: ``False``), objects with a ``for_json()``
+   method will use the return value of that method for encoding as JSON instead
+   of the object.
+
+   .. versionchanged:: 3.2.0
+     *for_json* is new in 3.2.0.
 
    .. method:: default(o)
 
diff --git a/simplejson/__init__.py b/simplejson/__init__.py
index 39f9fbf..a1095c5 100644
--- a/simplejson/__init__.py
+++ b/simplejson/__init__.py
@@ -142,6 +142,7 @@ _default_encoder = JSONEncoder(
     tuple_as_array=True,
     bigint_as_string=False,
     item_sort_key=None,
+    for_json=False,
 )
 
 def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
@@ -149,7 +150,7 @@ def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
         encoding='utf-8', default=None, use_decimal=True,
         namedtuple_as_object=True, tuple_as_array=True,
         bigint_as_string=False, sort_keys=False, item_sort_key=None,
-        **kw):
+        for_json=False, **kw):
     """Serialize ``obj`` as a JSON formatted stream to ``fp`` (a
     ``.write()``-supporting file-like object).
 
@@ -214,6 +215,10 @@ def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
     If *sort_keys* is true (default: ``False``), the output of dictionaries
     will be sorted by item.
 
+    If *for_json* is true (default: ``False``), objects with a ``for_json()``
+    method will use the return value of that method for encoding as JSON
+    instead of the object.
+
     To use a custom ``JSONEncoder`` subclass (e.g. one that overrides the
     ``.default()`` method to serialize additional types), specify it with
     the ``cls`` kwarg. NOTE: You should use *default* instead of subclassing
@@ -226,7 +231,8 @@ 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 item_sort_key and not kw):
+        and not bigint_as_string and not item_sort_key
+        and not for_json and not kw):
         iterable = _default_encoder.iterencode(obj)
     else:
         if cls is None:
@@ -240,6 +246,7 @@ def dump(obj, fp, skipkeys=False, ensure_ascii=True, check_circular=True,
             bigint_as_string=bigint_as_string,
             sort_keys=sort_keys,
             item_sort_key=item_sort_key,
+            for_json=for_json,
             **kw).iterencode(obj)
     # could accelerate with writelines in some versions of Python, at
     # a debuggability cost
@@ -252,7 +259,7 @@ def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
         encoding='utf-8', default=None, use_decimal=True,
         namedtuple_as_object=True, tuple_as_array=True,
         bigint_as_string=False, sort_keys=False, item_sort_key=None,
-        **kw):
+        for_json=False, **kw):
     """Serialize ``obj`` to a JSON formatted ``str``.
 
     If ``skipkeys`` is false then ``dict`` keys that are not basic types
@@ -312,6 +319,10 @@ def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
     If *sort_keys* is true (default: ``False``), the output of dictionaries
     will be sorted by item.
 
+    If *for_json* is true (default: ``False``), objects with a ``for_json()``
+    method will use the return value of that method for encoding as JSON
+    instead of the object.
+
     To use a custom ``JSONEncoder`` subclass (e.g. one that overrides the
     ``.default()`` method to serialize additional types), specify it with
     the ``cls`` kwarg. NOTE: You should use *default* instead of subclassing
@@ -325,7 +336,7 @@ def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
         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 sort_keys
-        and not item_sort_key and not kw):
+        and not item_sort_key and not for_json and not kw):
         return _default_encoder.encode(obj)
     if cls is None:
         cls = JSONEncoder
@@ -339,6 +350,7 @@ def dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True,
         bigint_as_string=bigint_as_string,
         sort_keys=sort_keys,
         item_sort_key=item_sort_key,
+        for_json=for_json,
         **kw).encode(obj)
 
 
diff --git a/simplejson/_speedups.c b/simplejson/_speedups.c
index ec4be4f..e08c2fd 100644
--- a/simplejson/_speedups.c
+++ b/simplejson/_speedups.c
@@ -170,6 +170,7 @@ typedef struct _PyEncoderObject {
     int bigint_as_string;
     PyObject *item_sort_key;
     PyObject *item_sort_kw;
+    int for_json;
 } PyEncoderObject;
 
 static PyMemberDef encoder_members[] = {
@@ -2587,22 +2588,22 @@ 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", "item_sort_key", "encoding", "Decimal", 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", "encoding", "Decimal", "for_json", NULL};
 
     PyEncoderObject *s;
     PyObject *markers, *defaultfn, *encoder, *indent, *key_separator;
     PyObject *item_separator, *sort_keys, *skipkeys, *allow_nan, *key_memo;
     PyObject *use_decimal, *namedtuple_as_object, *tuple_as_array;
-    PyObject *bigint_as_string, *item_sort_key, *encoding, *Decimal;
+    PyObject *bigint_as_string, *item_sort_key, *encoding, *Decimal, *for_json;
 
     assert(PyEncoder_Check(self));
     s = (PyEncoderObject *)self;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OOOOOOOOOOOOOOOOO:make_encoder", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OOOOOOOOOOOOOOOOOO: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,
-        &item_sort_key, &encoding, &Decimal))
+        &item_sort_key, &encoding, &Decimal, &for_json))
         return -1;
 
     s->markers = markers;
@@ -2654,6 +2655,7 @@ encoder_init(PyObject *self, PyObject *args, PyObject *kwds)
     s->sort_keys = sort_keys;
     s->item_sort_key = item_sort_key;
     s->Decimal = Decimal;
+    s->for_json = PyObject_IsTrue(for_json);
 
     Py_INCREF(s->markers);
     Py_INCREF(s->defaultfn);
@@ -2817,7 +2819,7 @@ encoder_listencode_obj(PyEncoderObject *s, JSON_Accu *rval, PyObject *obj, Py_ss
             if (encoded != NULL)
                 rv = _steal_accumulate(rval, encoded);
         }
-        else if (_has_for_json_hook(obj)) {
+        else if (s->for_json && _has_for_json_hook(obj)) {
             PyObject *newobj;
             if (Py_EnterRecursiveCall(" while encoding a JSON object"))
                 return rv;
diff --git a/simplejson/encoder.py b/simplejson/encoder.py
index 213fefb..9e95a22 100644
--- a/simplejson/encoder.py
+++ b/simplejson/encoder.py
@@ -121,7 +121,7 @@ class JSONEncoder(object):
             indent=None, separators=None, encoding='utf-8', default=None,
             use_decimal=True, namedtuple_as_object=True,
             tuple_as_array=True, bigint_as_string=False,
-            item_sort_key=None):
+            item_sort_key=None, for_json=False):
         """Constructor for JSONEncoder, with sensible defaults.
 
         If skipkeys is false, then it is a TypeError to attempt
@@ -183,6 +183,11 @@ class JSONEncoder(object):
         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.
+
+        If for_json is true (not the default), objects with a ``for_json()``
+        method will use the return value of that method for encoding as JSON
+        instead of the object.
+
         """
 
         self.skipkeys = skipkeys
@@ -195,6 +200,7 @@ class JSONEncoder(object):
         self.tuple_as_array = tuple_as_array
         self.bigint_as_string = bigint_as_string
         self.item_sort_key = item_sort_key
+        self.for_json = for_json
         if indent is not None and not isinstance(indent, string_types):
             indent = indent * ' '
         self.indent = indent
@@ -312,7 +318,7 @@ class JSONEncoder(object):
                 self.namedtuple_as_object, self.tuple_as_array,
                 self.bigint_as_string, self.item_sort_key,
                 self.encoding,
-                Decimal)
+                Decimal, self.for_json)
         else:
             _iterencode = _make_iterencode(
                 markers, self.default, _encoder, self.indent, floatstr,
@@ -320,7 +326,7 @@ class JSONEncoder(object):
                 self.skipkeys, _one_shot, self.use_decimal,
                 self.namedtuple_as_object, self.tuple_as_array,
                 self.bigint_as_string, self.item_sort_key,
-                self.encoding,
+                self.encoding, self.for_json,
                 Decimal=Decimal)
         try:
             return _iterencode(o, 0)
@@ -358,7 +364,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, _item_sort_key, _encoding,
+        _bigint_as_string, _item_sort_key, _encoding, _for_json,
         ## HACK: hand-optimized bytecode; turn globals into locals
         _PY3=PY3,
         ValueError=ValueError,
@@ -422,7 +428,10 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
                 yield buf + str(value)
             else:
                 yield buf
-                if isinstance(value, list):
+                for_json = _for_json and getattr(value, 'for_json', None)
+                if for_json and callable(for_json):
+                    chunks = _iterencode(for_json(), _current_indent_level)
+                elif isinstance(value, list):
                     chunks = _iterencode_list(value, _current_indent_level)
                 else:
                     _asdict = _namedtuple_as_object and getattr(value, '_asdict', None)
@@ -532,7 +541,10 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
             elif _use_decimal and isinstance(value, Decimal):
                 yield str(value)
             else:
-                if isinstance(value, list):
+                for_json = _for_json and getattr(value, 'for_json', None)
+                if for_json and callable(for_json):
+                    chunks = _iterencode(for_json(), _current_indent_level)
+                elif isinstance(value, list):
                     chunks = _iterencode_list(value, _current_indent_level)
                 else:
                     _asdict = _namedtuple_as_object and getattr(value, '_asdict', None)
@@ -571,32 +583,38 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
                        else ('"' + str(o) + '"'))
         elif isinstance(o, float):
             yield _floatstr(o)
-        elif isinstance(o, list):
-            for chunk in _iterencode_list(o, _current_indent_level):
-                yield chunk
         else:
-            _asdict = _namedtuple_as_object and getattr(o, '_asdict', None)
-            if _asdict and callable(_asdict):
-                for chunk in _iterencode_dict(_asdict(), _current_indent_level):
+            for_json = _for_json and getattr(o, 'for_json', None)
+            if for_json and callable(for_json):
+                for chunk in _iterencode(for_json(), _current_indent_level):
                     yield chunk
-            elif (_tuple_as_array and isinstance(o, tuple)):
+            elif isinstance(o, list):
                 for chunk in _iterencode_list(o, _current_indent_level):
                     yield chunk
-            elif isinstance(o, dict):
-                for chunk in _iterencode_dict(o, _current_indent_level):
-                    yield chunk
-            elif _use_decimal and isinstance(o, Decimal):
-                yield str(o)
             else:
-                if markers is not None:
-                    markerid = id(o)
-                    if markerid in markers:
-                        raise ValueError("Circular reference detected")
-                    markers[markerid] = o
-                o = _default(o)
-                for chunk in _iterencode(o, _current_indent_level):
-                    yield chunk
-                if markers is not None:
-                    del markers[markerid]
+                _asdict = _namedtuple_as_object and getattr(o, '_asdict', None)
+                if _asdict and callable(_asdict):
+                    for chunk in _iterencode_dict(_asdict(),
+                            _current_indent_level):
+                        yield chunk
+                elif (_tuple_as_array and isinstance(o, tuple)):
+                    for chunk in _iterencode_list(o, _current_indent_level):
+                        yield chunk
+                elif isinstance(o, dict):
+                    for chunk in _iterencode_dict(o, _current_indent_level):
+                        yield chunk
+                elif _use_decimal and isinstance(o, Decimal):
+                    yield str(o)
+                else:
+                    if markers is not None:
+                        markerid = id(o)
+                        if markerid in markers:
+                            raise ValueError("Circular reference detected")
+                        markers[markerid] = o
+                    o = _default(o)
+                    for chunk in _iterencode(o, _current_indent_level):
+                        yield chunk
+                    if markers is not None:
+                        del markers[markerid]
 
     return _iterencode
diff --git a/simplejson/tests/test_for_json.py b/simplejson/tests/test_for_json.py
new file mode 100644
index 0000000..2906bdb
--- /dev/null
+++ b/simplejson/tests/test_for_json.py
@@ -0,0 +1,126 @@
+import unittest
+import simplejson as json
+
+
+class ForJson(object):
+    def for_json(self):
+        return {'for_json': 1}
+
+
+class NestedForJson(object):
+    def for_json(self):
+        return {'nested': ForJson()}
+
+
+class ForJsonList(object):
+    def for_json(self):
+        return ['list']
+
+
+class DictForJson(dict):
+    def for_json(self):
+        return {'alpha': 1}
+
+
+class ListForJson(list):
+    def for_json(self):
+        return ['list']
+
+
+class TestForJsonWithSpeedups(unittest.TestCase):
+    def setUp(self):
+        if not json.encoder.c_make_encoder:
+            raise unittest.SkipTest("No speedups.")
+
+    @staticmethod
+    def _dump(obj):
+        return json.dumps(obj, for_json=True)
+
+    def test_for_json_encodes_stand_alone_object(self):
+        self.assertEqual(self._dump(ForJson()), '{"for_json": 1}')
+
+    def test_for_json_encodes_object_nested_in_dict(self):
+        self.assertEqual(self._dump({'hooray': ForJson()}), '{"hooray": '
+                '{"for_json": 1}}')
+
+    def test_for_json_encodes_object_nested_in_list_within_dict(self):
+        self.assertEqual(self._dump({'list': [0, ForJson(), 2, 3]}),
+                '{"list": [0, {"for_json": 1}, 2, 3]}')
+
+    def test_for_json_encodes_object_nested_within_object(self):
+        self.assertEqual(self._dump(NestedForJson()),
+                '{"nested": {"for_json": 1}}')
+
+    def test_for_json_encodes_list(self):
+        self.assertEqual(self._dump(ForJsonList()), '["list"]')
+
+    def test_for_json_encodes_list_within_object(self):
+        self.assertEqual(self._dump({'nested': ForJsonList()}),
+                '{"nested": ["list"]}')
+
+    def test_for_json_encodes_dict_subclass(self):
+        self.assertEqual(self._dump(DictForJson(a=1)), '{"alpha": 1}')
+
+    def test_for_json_encodes_list_subclass(self):
+        self.assertEqual(self._dump(ListForJson(['l'])), '["list"]')
+
+    def tset_for_json_ignored_if_not_true_with_dict_subclass(self):
+        self.assertEqual(json.dumps(DictForJson(a=1)), '{"a": 1}')
+
+    def test_for_json_ignored_if_not_true_with_list_subclass(self):
+        self.assertEqual(json.dumps(ListForJson(['l'])), '["l"]')
+
+    def test_raises_typeerror_if_for_json_not_true_with_object(self):
+        self.assertRaises(TypeError, json.dumps, ForJson())
+
+
+class TestForJsonWithoutSpeedups(unittest.TestCase):
+    def setUp(self):
+        if json.encoder.c_make_encoder:
+            json._toggle_speedups(False)
+
+    def tearDown(self):
+        if json.encoder.c_make_encoder:
+            json._toggle_speedups(True)
+
+    @staticmethod
+    def _dump(obj):
+        return json.dumps(obj, for_json=True)
+
+    def test_for_json_encodes_stand_alone_object(self):
+        self.assertEqual(self._dump(ForJson()), '{"for_json": 1}')
+
+    def test_for_json_encodes_object_nested_in_dict(self):
+        self.assertEqual(self._dump({'hooray': ForJson()}), '{"hooray": '
+                '{"for_json": 1}}')
+
+    def test_for_json_encodes_object_nested_in_list_within_dict(self):
+        self.assertEqual(self._dump({'list': [0, ForJson(), 2, 3]}),
+                '{"list": [0, {"for_json": 1}, 2, 3]}')
+
+    def test_for_json_encodes_object_nested_within_object(self):
+        self.assertEqual(self._dump(NestedForJson()),
+                '{"nested": {"for_json": 1}}')
+
+    def test_for_json_encodes_list(self):
+        self.assertEqual(self._dump(ForJsonList()), '["list"]')
+
+    def test_for_json_encodes_list_within_object(self):
+        self.assertEqual(self._dump({'nested': ForJsonList()}),
+                '{"nested": ["list"]}')
+
+    def test_for_json_encodes_dict_subclass(self):
+        self.assertEqual(self._dump(DictForJson(a=1)), '{"alpha": 1}')
+
+    def test_for_json_encodes_list_subclass(self):
+        self.assertEqual(self._dump(ListForJson(['l'])), '["list"]')
+
+    def tset_for_json_ignored_if_not_true_with_dict_subclass(self):
+        self.assertEqual(json.dumps(DictForJson(a=1)), '{"a": 1}')
+
+    def test_for_json_ignored_if_not_true_with_list_subclass(self):
+        self.assertEqual(json.dumps(ListForJson(['l'])), '["l"]')
+
+    def test_raises_typeerror_if_for_json_not_true_with_object(self):
+        self.assertRaises(TypeError, json.dumps, ForJson())
+