Support custom dict key types in JSON

msgspec/_core.c

@@ -13185,11 +13185,20 @@ json_encode_set(EncoderState *self, PyObject *obj)
     return status;
 }
 
+static int json_encode_dict_key_noinline(EncoderState *, PyObject *);
+
+static MS_INLINE int
+json_encode_dict_key(EncoderState *self, PyObject *key) {
+    if (MS_LIKELY(PyUnicode_Check(key))) {
+        return json_encode_str(self, key);
+    }
+    return json_encode_dict_key_noinline(self, key);
+}
+
 static MS_NOINLINE int
-json_encode_dict_key(EncoderState *self, PyObject *obj) {
+json_encode_dict_key_noinline(EncoderState *self, PyObject *obj) {
     PyTypeObject *type = Py_TYPE(obj);
 
-    /* PyUnicode_Type is handled inline in json_encode_dict */
     if (type == &PyLong_Type) {
         return json_encode_long_as_str(self, obj);
     }
@@ -13220,6 +13229,18 @@ json_encode_dict_key(EncoderState *self, PyObject *obj) {
     else if (PyType_IsSubtype(type, (PyTypeObject *)(self->mod->UUIDType))) {
         return json_encode_uuid(self, obj);
     }
+    else if (self->enc_hook != NULL) {
+        int status = -1;
+        PyObject *temp;
+        temp = CALL_ONE_ARG(self->enc_hook, obj);
+        if (temp == NULL) return -1;
+        if (!Py_EnterRecursiveCall(" while serializing an object")) {
+            status = json_encode_dict_key(self, temp);
+            Py_LeaveRecursiveCall();
+        }
+        Py_DECREF(temp);
+        return status;
+    }
     else {
         PyErr_SetString(
             PyExc_TypeError,
@@ -13241,12 +13262,7 @@ json_encode_dict(EncoderState *self, PyObject *obj)
     if (ms_write(self, "{", 1) < 0) return -1;
     if (Py_EnterRecursiveCall(" while serializing an object")) return -1;
     while (PyDict_Next(obj, &pos, &key, &val)) {
-        if (MS_LIKELY(PyUnicode_Check(key))) {
-            if (json_encode_str(self, key) < 0) goto cleanup;
-        }
-        else {
-            if (json_encode_dict_key(self, key) < 0) goto cleanup;
-        }
+        if (json_encode_dict_key(self, key) < 0) goto cleanup;
         if (ms_write(self, ":", 1) < 0) goto cleanup;
         if (json_encode_inline(self, val) < 0) goto cleanup;
         if (ms_write(self, ",", 1) < 0) goto cleanup;
@@ -16513,6 +16529,9 @@ json_decode_dict_key_fallback(
         else {
             out = PyUnicode_DecodeUTF8(view, size, NULL);
         }
+        if (MS_UNLIKELY(type->types & (MS_TYPE_CUSTOM | MS_TYPE_CUSTOM_GENERIC))) {
+            return ms_decode_custom(out, self->dec_hook, type, path);
+        }
         return ms_check_str_constraints(out, type, path);
     }
     if (type->types & (

tests/test_json.py

@@ -2096,6 +2096,46 @@ class mystr(str):
         msg = msgspec.json.encode({mystr("test"): 1})
         assert msg == b'{"test":1}'
 
+    def test_encode_dict_custom_key(self):
+        class Custom:
+            def __init__(self, value):
+                self.value = value
+
+        msg = msgspec.json.encode(
+            {Custom("a"): 1, Custom("b"): 2}, enc_hook=lambda x: x.value
+        )
+        assert msg == b'{"a":1,"b":2}'
+
+        def enc_hook(x):
+            raise TypeError("Oh no!")
+
+        with pytest.raises(TypeError):
+            msgspec.json.encode({Custom("x"): 1}, enc_hook=enc_hook)
+
+        with pytest.raises(RecursionError):
+            msgspec.json.encode({Custom("x"): 1}, enc_hook=lambda x: x)
+
+    def test_decode_dict_custom_key(self):
+        class Custom:
+            def __init__(self, value):
+                self.value = value
+
+            def __hash__(self):
+                return hash(self.value)
+
+            def __eq__(self, other):
+                return self.value == other.value
+
+        def dec_hook(typ, obj):
+            if typ == Custom:
+                return Custom(obj)
+            raise NotImplementedError
+
+        msg = b'{"a":1,"b":2}'
+
+        obj = msgspec.json.decode(msg, type=Dict[Custom, int], dec_hook=dec_hook)
+        assert obj == {Custom("a"): 1, Custom("b"): 2}
+
     @pytest.mark.parametrize(
         "s, error",
         [