Преглед изворни кода

Support fallback to Python's json module if orjson not available

Rino Beeli пре 2 година
родитељ
комит
e1b3bac74f

+ 20 - 4
nicegui/json/__init__.py

@@ -1,15 +1,31 @@
 """
 Custom json module. Provides dumps and loads implementations
-wrapping the `orjson` package.
+wrapping the orjson package. If the orjson package is not available,
+the standard Python json module is used.
 
-Custom module required in order to override json-module used
+This custom module is required in order to override the json-module used
 in socketio.AsyncServer, which expects a module as parameter
 to override Python's default json module.
 """
 
-from nicegui.json.orjson_wrapper import dumps, loads
+try:
+    # orjson not available on all platforms, fallback to Python's json module if not available
+    import orjson
+    raise ImportError()
+    has_orjson = True
+except:
+    has_orjson = False
+    print('Using Python json')
+
+
+if has_orjson:
+    from nicegui.json.orjson_wrapper import NiceGUIJSONResponse, dumps, loads
+else:
+    from nicegui.json.builtin_wrapper import NiceGUIJSONResponse, dumps, loads
+
 
 __all__ = [
     'dumps',
-    'loads'
+    'loads',
+    'NiceGUIJSONResponse'
 ]

+ 52 - 0
nicegui/json/builtin_wrapper.py

@@ -0,0 +1,52 @@
+import json
+from typing import Any, Optional, Tuple
+
+import numpy as np
+from fastapi import Response
+
+
+def dumps(obj: Any, sort_keys: bool = False, separators: Optional[Tuple[str, str]] = None):
+    """Serializes a Python object to a JSON-encoded string.
+
+    This implementation uses Python's default json module, but extends it
+    in order to support numpy arrays.
+    """
+    if separators is None:
+        separators = (',', ':')
+    return json.dumps(
+        obj,
+        sort_keys=sort_keys,
+        separators=separators,
+        indent=None,
+        allow_nan=True,
+        ensure_ascii=False,
+        cls=NumpyJsonEncoder)
+
+
+def loads(value: str) -> Any:
+    """Deserialize a JSON-encoded string to a corresponding Python object/value.
+
+    Uses Python's default json module internally.
+    """
+    return json.loads(value)
+
+
+class NiceGUIJSONResponse(Response):
+    """FastAPI response class to support our custom json serializer implementation."""
+    media_type = 'application/json'
+
+    def render(self, content: Any) -> bytes:
+        return dumps(content).encode("utf-8")
+
+
+class NumpyJsonEncoder(json.JSONEncoder):
+    """Special json encoder that supporty numpy arrays."""
+
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return json.JSONEncoder.default(self, obj)

+ 0 - 13
nicegui/json/fastapi.py

@@ -1,13 +0,0 @@
-from typing import Any
-
-import orjson
-from fastapi import Response
-
-from nicegui.json.orjson_wrapper import ORJSON_OPTS
-
-
-class NiceGUIJSONResponse(Response):
-    media_type = 'application/json'
-
-    def render(self, content: Any) -> bytes:
-        return orjson.dumps(content, option=ORJSON_OPTS)

+ 9 - 0
nicegui/json/orjson_wrapper.py

@@ -1,6 +1,7 @@
 from typing import Any, Optional, Tuple
 
 import orjson
+from fastapi import Response
 
 ORJSON_OPTS = orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS
 
@@ -34,3 +35,11 @@ def loads(value: str) -> Any:
     Uses package `orjson` internally.
     """
     return orjson.loads(value)
+
+
+class NiceGUIJSONResponse(Response):
+    """FastAPI response class to support our custom json serializer implementation."""
+    media_type = 'application/json'
+
+    def render(self, content: Any) -> bytes:
+        return orjson.dumps(content, option=ORJSON_OPTS)

+ 1 - 1
nicegui/nicegui.py

@@ -11,7 +11,7 @@ from fastapi.staticfiles import StaticFiles
 from fastapi_socketio import SocketManager
 
 from nicegui import json
-from nicegui.json.fastapi import NiceGUIJSONResponse
+from nicegui.json import NiceGUIJSONResponse
 
 from . import background_tasks, binding, globals, outbox
 from .app import App

+ 3 - 1
tests/test_plotly.py

@@ -1,3 +1,4 @@
+import numpy as np
 import plotly.graph_objects as go
 
 from nicegui import ui
@@ -11,7 +12,8 @@ def test_plotly(screen: Screen):
     plot = ui.plotly(fig)
 
     ui.button('Add trace', on_click=lambda: (
-        fig.add_trace(go.Scatter(x=[0, 1, 2], y=[2, 1, 0], name='Trace 2')),
+        # test numpy array support for value arrays
+        fig.add_trace(go.Scatter(x=np.array([0, 1, 2]), y=np.array([2, 1, 0]), name='Trace 2')),
         plot.update()
     ))