Add fallback function to encode custom Python objects

cjson.c

@@ -15,12 +15,12 @@ typedef struct JSONData {
     int  all_unicode; // make all output strings unicode if true
 } JSONData;
 
-static PyObject* encode_object(PyObject *object);
+static PyObject* encode_object(PyObject *object, PyObject *fallback);
 static PyObject* encode_string(PyObject *object);
 static PyObject* encode_unicode(PyObject *object);
-static PyObject* encode_tuple(PyObject *object);
-static PyObject* encode_list(PyObject *object);
-static PyObject* encode_dict(PyObject *object);
+static PyObject* encode_tuple(PyObject *object, PyObject *fallback);
+static PyObject* encode_list(PyObject *object, PyObject *fallback);
+static PyObject* encode_dict(PyObject *object, PyObject *fallback);
 
 static PyObject* decode_json(JSONData *jsondata);
 static PyObject* decode_null(JSONData *jsondata);
@@ -799,7 +799,7 @@ encode_unicode(PyObject *unicode)
  */
 
 static PyObject*
-encode_tuple(PyObject *tuple)
+encode_tuple(PyObject *tuple, PyObject *fallback)
 {
     Py_ssize_t i, n;
     PyObject *s, *temp;
@@ -816,7 +816,7 @@ encode_tuple(PyObject *tuple)
 
     /* Do repr() on each element. */
     for (i = 0; i < n; ++i) {
-        s = encode_object(v->ob_item[i]);
+        s = encode_object(v->ob_item[i], fallback);
         if (s == NULL)
             goto Done;
         PyTuple_SET_ITEM(pieces, i, s);
@@ -864,7 +864,7 @@ encode_tuple(PyObject *tuple)
  *   represented in JSON.
  */
 static PyObject*
-encode_list(PyObject *list)
+encode_list(PyObject *list, PyObject *fallback)
 {
     Py_ssize_t i;
     PyObject *s, *temp;
@@ -893,7 +893,7 @@ encode_list(PyObject *list)
      * so must refetch the list size on each iteration. */
     for (i = 0; i < v->ob_size; ++i) {
         int status;
-        s = encode_object(v->ob_item[i]);
+        s = encode_object(v->ob_item[i], fallback);
         if (s == NULL)
             goto Done;
         status = PyList_Append(pieces, s);
@@ -947,7 +947,7 @@ encode_list(PyObject *list)
  *   be represented in JSON.
  */
 static PyObject*
-encode_dict(PyObject *dict)
+encode_dict(PyObject *dict, PyObject *fallback)
 {
     Py_ssize_t i;
     PyObject *s, *temp, *colon = NULL;
@@ -991,9 +991,9 @@ encode_dict(PyObject *dict)
 
         /* Prevent repr from deleting value during key format. */
         Py_INCREF(value);
-        s = encode_object(key);
+        s = encode_object(key, fallback);
         PyString_Concat(&s, colon);
-        PyString_ConcatAndDel(&s, encode_object(value));
+        PyString_ConcatAndDel(&s, encode_object(value, fallback));
         Py_DECREF(value);
         if (s == NULL)
             goto Done;
@@ -1039,7 +1039,7 @@ encode_dict(PyObject *dict)
 
 
 static PyObject*
-encode_object(PyObject *object)
+encode_object(PyObject *object, PyObject *fallback)
 {
     if (object == Py_True) {
         return PyString_FromString("true");
@@ -1070,25 +1070,39 @@ encode_object(PyObject *object)
         PyObject *result;
         if (Py_EnterRecursiveCall(" while encoding a JSON array from a Python list"))
             return NULL;
-        result = encode_list(object);
+        result = encode_list(object, fallback);
         Py_LeaveRecursiveCall();
         return result;
     } else if (PyTuple_Check(object)) {
         PyObject *result;
         if (Py_EnterRecursiveCall(" while encoding a JSON array from a Python tuple"))
             return NULL;
-        result = encode_tuple(object);
+        result = encode_tuple(object, fallback);
         Py_LeaveRecursiveCall();
         return result;
     } else if (PyDict_Check(object)) { // use PyMapping_Check(object) instead? -Dan
         PyObject *result;
         if (Py_EnterRecursiveCall(" while encoding a JSON object"))
             return NULL;
-        result = encode_dict(object);
+        result = encode_dict(object, fallback);
+        Py_LeaveRecursiveCall();
+        return result;
+    } else if (fallback) {
+        PyObject *args, *resolve, *result;
+        if (Py_EnterRecursiveCall(" while encoding a non-primitive Python object"))
+            return NULL;
+        args = PyTuple_Pack(1, object);
+        resolve = PyObject_CallObject(fallback, args);
+        Py_DECREF(args);
+        result = PyErr_Occurred() ? NULL : encode_object(resolve, fallback);
+        Py_XDECREF(resolve);
         Py_LeaveRecursiveCall();
         return result;
     } else {
-        PyErr_SetString(JSON_EncodeError, "object is not JSON encodable");
+        PyObject *repr = PyObject_Repr(object);
+        PyErr_Format(JSON_EncodeError, "object %s is not JSON encodable",
+            PyString_AsString(repr));
+        Py_DECREF(repr);
         return NULL;
     }
 }
@@ -1097,9 +1111,21 @@ encode_object(PyObject *object)
 /* Encode object into its JSON representation */
 
 static PyObject*
-JSON_encode(PyObject *self, PyObject *object)
+JSON_encode(PyObject *self, PyObject *args, PyObject *kwargs)
 {
-    return encode_object(object);
+    static char *kwlist[] = {"obj", "default", NULL};
+    PyObject *object, *fallback = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O:decode",
+        kwlist, &object, &fallback
+    )) return NULL;
+    if (fallback && PyObject_Not(fallback)) fallback = NULL;
+    if (fallback && !PyCallable_Check(fallback)) {
+        PyErr_Format(PyExc_ValueError,
+            "The 'default' argument %s is not callable",
+            PyObject_Repr(fallback));
+        return NULL;
+    };
+    return encode_object(object, fallback);
 }
 
 
@@ -1159,8 +1185,11 @@ JSON_decode(PyObject *self, PyObject *args, PyObject *kwargs)
 /* List of functions defined in the module */
 
 static PyMethodDef cjson_methods[] = {
-    {"encode", (PyCFunction)JSON_encode,  METH_O,
-    PyDoc_STR("encode(object) -> generate the JSON representation for object.")},
+    {"encode", (PyCFunction)JSON_encode,  METH_VARARGS|METH_KEYWORDS,
+    PyDoc_STR("encode(object, default=Null) -> generate the JSON representation for object.\n"
+              "The optional argument `default' is function that gets called for objects\n"
+              "that can’t otherwise be serialized. It should return a JSON encodable\n"
+              "version of the object or raise cjson.EncodeError.")},
 
     {"decode", (PyCFunction)JSON_decode,  METH_VARARGS|METH_KEYWORDS,
     PyDoc_STR("decode(string, all_unicode=False) -> parse the JSON representation into\n"

jsontest.py

@@ -20,6 +20,9 @@
 ## License along with this library; if not, write to the Free Software
 ## Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 
+from datetime import datetime
+from decimal import Decimal
+
 import unittest
 
 import cjson
@@ -141,7 +144,7 @@ def testReadBadArray(self):
 
     def doReadBadArray(self):
         cjson.decode('[1,2,3,,]')
-        
+
     def testReadBadObjectSyntax(self):
         self.assertRaises(_exception, self.doReadBadObjectSyntax)
 
@@ -159,7 +162,7 @@ def testReadIntegerValue(self):
     def testReadNegativeIntegerValue(self):
         obj = cjson.decode('{ "key" : -44 }')
         self.assertEqual({ "key" : -44 }, obj)
-        
+
     def testReadFloatValue(self):
         obj = cjson.decode('{ "age" : 44.5 }')
         self.assertEqual({ "age" : 44.5 }, obj)
@@ -176,7 +179,7 @@ def doReadBadNumber(self):
 
     def testReadSmallObject(self):
         obj = cjson.decode('{ "name" : "Patrick", "age":44} ')
-        self.assertEqual({ "age" : 44, "name" : "Patrick" }, obj)        
+        self.assertEqual({ "age" : 44, "name" : "Patrick" }, obj)
 
     def testReadEmptyArray(self):
         obj = cjson.decode('[]')
@@ -331,7 +334,18 @@ def testWriteLongUnicode(self):
                          u'\u1234\u1234\u1234\u1234\u1234\u1234')
         self.assertEqual(r'"\ud834\udd1e\ud834\udd1e\ud834\udd1e\ud834\udd1e'
                          r'\u1234\u1234\u1234\u1234\u1234\u1234"', s)
-
+
+    def testWriteCustomObject(self):
+        def fallback(obj):
+            if isinstance(obj, Decimal):
+                return float(obj)
+            raise cjson.EncodeError(obj)
+        with self.assertRaises(cjson.EncodeError):
+            cjson.encode(Decimal(1.23))
+        self.assertEqual(cjson.encode(Decimal(1.23), fallback), '1.23')
+        with self.assertRaises(cjson.EncodeError):
+            cjson.encode(datetime.now(), fallback)
+
 def main():
     unittest.main()