Abstract orjson for message.py.

Author: adamghill

django_unicorn/message.py

@@ -2,11 +2,9 @@
 
 from django.http.response import HttpResponseRedirect
 
-import orjson
-
 from .components import HashUpdate, LocationUpdate, PollUpdate
 from .errors import UnicornViewError
-from .serializer import dumps
+from .serializer import JSONDecodeError, dumps, loads
 from .utils import generate_checksum
 
 
@@ -22,9 +20,9 @@ def __init__(self, request):
         self.body = {}
 
         try:
-            self.body = orjson.loads(request.body)
+            self.body = loads(request.body)
             assert self.body, "Invalid JSON body"
-        except orjson.JSONDecodeError as e:
+        except JSONDecodeError as e:
             raise UnicornViewError("Body could not be parsed") from e
 
         self.data = self.body.get("data")
@@ -49,7 +47,7 @@ def validate_checksum(self):
         checksum = self.body.get("checksum")
         assert checksum, "Missing checksum"
 
-        generated_checksum = generate_checksum(orjson.dumps(self.data))
+        generated_checksum = generate_checksum(dumps(self.data, fix_floats=False))
         assert checksum == generated_checksum, "Checksum does not match"
 
 
@@ -93,8 +91,8 @@ def value(self, value):
 
     def get_data(self):
         try:
-            serialized_value = orjson.loads(dumps(self.value))
-            serialized_params = orjson.loads(dumps(self.params))
+            serialized_value = loads(dumps(self.value))
+            serialized_params = loads(dumps(self.params))
 
             return {
                 "method": self.method_name,

django_unicorn/serializer.py

@@ -13,6 +13,10 @@
 import orjson
 
 
+class JSONDecodeError(Exception):
+    pass
+
+
 def _get_model_dict(obj: Any) -> dict:
     """
     Serializes Django models. Uses the built-in Django JSON serializer, but moves the data around to
@@ -108,22 +112,32 @@ def _dumps(serialized_data):
     return dumped_data
 
 
-def dumps(data: dict) -> str:
+def dumps(data: dict, fix_floats: bool = True) -> str:
     """
     Converts the passed-in dictionary to a string representation.
 
     Handles the following objects: dataclass, datetime, enum, float, int, numpy, str, uuid,
     Django Model, Django QuerySet, any object with `to_json` method.
+
+    Args:
+        param fix_floats: Whether any floats should be converted to strings. Defaults to `True`,
+            but will be faster without it.
     """
 
     serialized_data = orjson.dumps(data, default=_json_serializer)
 
-    return _dumps(serialized_data)
+    if fix_floats:
+        return _dumps(serialized_data)
+
+    return serialized_data
 
 
-def loads(obj: str) -> dict:
+def loads(str: str) -> dict:
     """
     Converts a string representation to dictionary.
     """
 
-    return orjson.loads(obj)
+    try:
+        return orjson.loads(str)
+    except orjson.JSONDecodeError as e:
+        raise JSONDecodeError from e