Support for a default hook.

Author: TkTech

tests/test_serialization.py

@@ -13,9 +13,23 @@ def test_serialize_unknown():
     object we can't handle.
     """
     with pytest.raises(TypeError):
-        yyjson.Document({
-            "example": ClassThatCantBeSerialized()
-        })
+        yyjson.Document({"example": ClassThatCantBeSerialized()})
 
     with pytest.raises(TypeError):
         yyjson.Document([ClassThatCantBeSerialized()])
+
+
+def test_serialize_default_func():
+    """
+    Ensure that we can serialize an object by providing a default function.
+    """
+
+    def default(obj):
+        if isinstance(obj, ClassThatCantBeSerialized):
+            return "I'm a string now!"
+        raise TypeError(f"Can't serialize {obj}")
+
+    doc = yyjson.Document(
+        {"example": ClassThatCantBeSerialized()}, default=default
+    )
+    assert doc.as_obj["example"] == "I'm a string now!"

yyjson/__init__.pyi

@@ -1,5 +1,6 @@
 import enum
-from typing import Any, Optional, List, Dict, Union
+from pathlib import Path
+from typing import Any, Optional, List, Dict, Union, Callable
 
 class ReaderFlags(enum.IntFlag):
     STOP_WHEN_DONE = 0x02
@@ -15,15 +16,22 @@ class WriterFlags(enum.IntFlag):
     ALLOW_INF_AND_NAN = 0x08
     INF_AND_NAN_AS_NULL = 0x10
 
-Content = Union[str, bytes, List, Dict]
+Content = Union[str, bytes, List, Dict, Path]
 
 class Document:
     as_obj: Any
-    def __init__(self, content: Content, flags: Optional[ReaderFlags] = ...): ...
+    def __init__(
+        self,
+        content: Content,
+        flags: Optional[ReaderFlags] = ...,
+        default: Callable[[Any], Any] = ...,
+    ): ...
     def __len__(self) -> int: ...
     def get_pointer(self, pointer: str) -> Any: ...
     def dumps(
-        self, flags: Optional[WriterFlags] = ..., at_pointer: Optional[str] = ...
+        self,
+        flags: Optional[WriterFlags] = ...,
+        at_pointer: Optional[str] = ...,
     ) -> str: ...
     def patch(
         self,

yyjson/document.c

@@ -283,10 +283,22 @@ PyTypeObject *type_for_conversion(PyObject *obj) {
  * Recursively convert a Python object into yyjson elements.
  */
 static inline yyjson_mut_val *mut_primitive_to_element(
-    yyjson_mut_doc *doc, PyObject *obj
+    DocumentObject *self,
+    yyjson_mut_doc *doc,
+    PyObject *obj
 ) {
   const PyTypeObject *ob_type = type_for_conversion(obj);
 
+  if (yyjson_unlikely(ob_type == NULL) && self->default_func != NULL) {
+    PyObject *result = PyObject_CallFunctionObjArgs(self->default_func, obj, NULL);
+    if (result == NULL) {
+      return NULL;
+    }
+    yyjson_mut_val *val = mut_primitive_to_element(self, doc, result);
+    Py_DECREF(result);
+    return val;
+  }
+
   if (ob_type == &PyUnicode_Type) {
     Py_ssize_t str_len;
     const char *str = PyUnicode_AsUTF8AndSize(obj, &str_len);
@@ -321,7 +333,7 @@ static inline yyjson_mut_val *mut_primitive_to_element(
     yyjson_mut_val *val = yyjson_mut_arr(doc);
     yyjson_mut_val *object_value = NULL;
     for (Py_ssize_t i = 0; i < PyList_GET_SIZE(obj); i++) {
-      object_value = mut_primitive_to_element(doc, PyList_GET_ITEM(obj, i));
+      object_value = mut_primitive_to_element(self, doc, PyList_GET_ITEM(obj, i));
 
       if (yyjson_unlikely(object_value == NULL)) {
         return NULL;
@@ -339,7 +351,7 @@ static inline yyjson_mut_val *mut_primitive_to_element(
     while (PyDict_Next(obj, &i, &key, &value)) {
       Py_ssize_t str_len;
       const char *str = PyUnicode_AsUTF8AndSize(key, &str_len);
-      object_value = mut_primitive_to_element(doc, value);
+      object_value = mut_primitive_to_element(self, doc, value);
       if (yyjson_unlikely(object_value == NULL)) {
         return NULL;
       }
@@ -359,11 +371,9 @@ static inline yyjson_mut_val *mut_primitive_to_element(
   } else if (obj == Py_None) {
     return yyjson_mut_null(doc);
   } else {
-    PyErr_SetString(
-        PyExc_TypeError,
-        // TODO: We can provide a much better error here. Also add support
-        // for a default hook.
-        "Tried to serialize an object we don't know how to handle."
+    PyErr_Format(PyExc_TypeError,
+      "Object of type '%s' is not JSON serializable",
+      Py_TYPE(obj)->tp_name
     );
     return NULL;
   }
@@ -372,6 +382,7 @@ static inline yyjson_mut_val *mut_primitive_to_element(
 static void Document_dealloc(DocumentObject *self) {
   if (self->i_doc != NULL) yyjson_doc_free(self->i_doc);
   if (self->m_doc != NULL) yyjson_mut_doc_free(self->m_doc);
+  Py_XDECREF(self->default_func);
   Py_TYPE(self)->tp_free((PyObject *)self);
 }
 
@@ -426,20 +437,33 @@ PyDoc_STRVAR(
     ":param content: The initial content of the document.\n"
     ":type content: ``str``, ``bytes``, ``Path``, ``dict``, ``list``\n"
     ":param flags: Flags that modify the document parsing behaviour.\n"
-    ":type flags: :class:`ReaderFlags`, optional"
+    ":type flags: :class:`ReaderFlags`, optional\n"
+    ":param default: A function called to convert objects that are not\n"
+    "                JSON serializable. Should return a JSON serializable version\n"
+    "                of the object or raise a TypeError.\n"
+    ":type default: callable, optional"
 );
 static int Document_init(DocumentObject *self, PyObject *args, PyObject *kwds) {
-  static char *kwlist[] = {"content", "flags", NULL};
+  static char *kwlist[] = {"content", "flags", "default", NULL};
   PyObject *content;
+  PyObject *default_func = NULL;
   yyjson_read_err err;
   yyjson_read_flag r_flag = 0;
 
   if (!PyArg_ParseTupleAndKeywords(
-          args, kwds, "O|$I", kwlist, &content, &r_flag
+          args, kwds, "O|$IO", kwlist, &content, &r_flag, &default_func
       )) {
     return -1;
   }
 
+  if (default_func && !PyCallable_Check(default_func)) {
+    PyErr_SetString(PyExc_TypeError, "default must be callable");
+    return -1;
+  }
+
+  self->default_func = default_func;
+  Py_XINCREF(default_func);
+
   if (yyjson_unlikely(pathlib == NULL)) {
     pathlib = PyImport_ImportModule("pathlib");
     if (yyjson_unlikely(pathlib == NULL)) {
@@ -523,7 +547,7 @@ static int Document_init(DocumentObject *self, PyObject *args, PyObject *kwds) {
       return -1;
     }
 
-    yyjson_mut_val *val = mut_primitive_to_element(self->m_doc, content);
+    yyjson_mut_val *val = mut_primitive_to_element(self, self->m_doc, content);
 
     if (val == NULL) {
       return -1;

yyjson/document.h

@@ -17,6 +17,8 @@ typedef struct {
   yyjson_doc* i_doc;
   /** The memory allocator in use for this document. */
   yyjson_alc* alc;
+  /** default callback for serializing unknown types. */
+  PyObject* default_func;
 } DocumentObject;
 
 extern PyTypeObject DocumentType;