Merge pull request #2 from wlatanowicz/fix-compatibility

Compatibility with json.dumps on bool, None and invalid keys

json_stream_generator/__init__.py

@@ -9,6 +9,13 @@
 __author__ = "Wiktor Latanowicz"
 
 
+KEY_CONVERTION_LUT = {
+    None: "null",
+    True: "true",
+    False: "false",
+}
+
+
 def json_generator(obj: Any, depth: int = 1) -> Generator[str, None, None]:
     if depth < 1:
         yield json.dumps(obj)
@@ -17,7 +24,13 @@ def json_generator(obj: Any, depth: int = 1) -> Generator[str, None, None]:
         for i, (key, value) in enumerate(obj.items()):
             if i > 0:
                 yield ", "
-            yield json.dumps(str(key)) + ": "
+
+            if not isinstance(key, (str, int, float, bool)) and key is not None:
+                raise TypeError(
+                    f"keys must be str, int, float, bool or None, not {key.__class__.__name__}"
+                )
+
+            yield json.dumps(str(KEY_CONVERTION_LUT.get(key, key))) + ": "
             yield from json_generator(value, depth=depth - 1)
         yield "}"
     elif isinstance(obj, (list, tuple)) or inspect.isgenerator(obj):

json_stream_generator/rest_framework/mixins.py

@@ -1,8 +1,7 @@
+import rest_framework
 from django.db.models.query import QuerySet
 from django.http import StreamingHttpResponse
 
-import rest_framework
-
 from .. import json_generator
 
 

tests/test_json_generator.py

@@ -1,4 +1,5 @@
 import json
+from datetime import datetime
 
 import pytest
 
@@ -77,3 +78,36 @@ def test_simple_values():
     result = "".join(json_generator(obj))
 
     assert expected == result
+
+
+def test_different_key_types():
+    obj = {
+        123: "A",
+        123.456: "B",
+        True: "C",
+        False: "D",
+        None: "E",
+        "abc": "F",
+    }
+    expected = json.dumps(obj)
+    result = "".join(json_generator(obj))
+
+    assert expected == result
+
+
+# fmt: off
+@pytest.mark.parametrize(
+    "obj",
+    [
+        {(1, 2): "A"},
+        [{(1, 2): "A"}],
+        {datetime.now(): "A"},
+        [{datetime.now(): "A"}],
+        ({datetime.now(): "A"} for _ in range(5)),
+        ({(1, 2): "A"} for _ in range(5)),
+    ]
+)
+# fmt: on
+def test_invalid_keys(obj):
+    with pytest.raises(TypeError):
+        "".join(json_generator(obj))