Fix ref counting on non-string dict keys

For bytes, there was an extraneous INCREF; PyIter_Next returns a new reference. For other non-strings, the original itemName before converting to a string was never dereferenced. Fixes #419
ultrajson · Apr 7, 2022 · 62dec8d · 62dec8d
1 parent 2d1f088
commit 62dec8d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/python/objToJSON.c b/python/objToJSON.c
@@ -254,15 +254,13 @@ static int Dict_iterNext(JSOBJ obj, JSONTypeContext *tc)
       return 1;
     }
 
+    itemNameTmp = GET_TC(tc)->itemName;
     GET_TC(tc)->itemName = PyObject_Str(GET_TC(tc)->itemName);
+    Py_DECREF(itemNameTmp);
     itemNameTmp = GET_TC(tc)->itemName;
     GET_TC(tc)->itemName = PyUnicode_AsUTF8String (GET_TC(tc)->itemName);
     Py_DECREF(itemNameTmp);
   }
-  else
-  {
-    Py_INCREF(GET_TC(tc)->itemName);
-  }
   PRINTMARK();
   return 1;
 }

diff --git a/tests/test_ujson.py b/tests/test_ujson.py
@@ -239,11 +239,11 @@ def test_encode_dict_values_ref_counting():
 @pytest.mark.skipif(
     hasattr(sys, "pypy_version_info"), reason="PyPy uses incompatible GC"
 )
-def test_encode_dict_key_ref_counting():
+@pytest.mark.parametrize("key", ["key", b"key", 1, True, None])
+def test_encode_dict_key_ref_counting(key):
     import gc
 
     gc.collect()
-    key = "key"
     data = {key: "abc"}
     ref_count = sys.getrefcount(key)
     ujson.dumps(data)