浏览代码

Extend json serializer to support numpy arrays of dtype=object

Rino Beeli 2 年之前
父节点
当前提交
0f6504d7fa
共有 2 个文件被更改,包括 10 次插入2 次删除
  1. 9 2
      nicegui/json/orjson_wrapper.py
  2. 1 0
      tests/test_json.py

+ 9 - 2
nicegui/json/orjson_wrapper.py

@@ -1,5 +1,6 @@
 from typing import Any, Optional, Tuple
 
+import numpy as np
 import orjson
 from fastapi import Response
 
@@ -26,7 +27,7 @@ def dumps(obj: Any, sort_keys: bool = False, separators: Optional[Tuple[str, str
     if sort_keys:
         opts |= orjson.OPT_SORT_KEYS
 
-    return orjson.dumps(obj, option=opts).decode('utf-8')
+    return orjson.dumps(obj, option=opts, default=_orjson_converter).decode('utf-8')
 
 
 def loads(value: str) -> Any:
@@ -37,6 +38,12 @@ def loads(value: str) -> Any:
     return orjson.loads(value)
 
 
+def _orjson_converter(obj):
+    """Custom serializer/converter, e.g. for numpy object arrays."""
+    if isinstance(obj, np.ndarray) and obj.dtype == np.object_:
+        return obj.tolist()
+
+
 class NiceGUIJSONResponse(Response):
     """FastAPI response class to support our custom json serializer implementation.
 
@@ -45,4 +52,4 @@ class NiceGUIJSONResponse(Response):
     media_type = 'application/json'
 
     def render(self, content: Any) -> bytes:
-        return orjson.dumps(content, option=ORJSON_OPTS)
+        return orjson.dumps(content, option=ORJSON_OPTS, default=_orjson_converter)

+ 1 - 0
tests/test_json.py

@@ -43,6 +43,7 @@ def test_json():
         np.array([1.0, 0]),
         np.array([0, False, np.pi]),
         np.array(['2010-10-17 07:15:30', '2011-05-13 08:20:35', '2013-01-15 09:09:09'], dtype=np.datetime64),
+        np.array([1.0, None, "test"], dtype=np.object_)
     ]
 
     for test in tests: