diff --git a/src/py/kaleido/_kaleido_tab/_tab.py b/src/py/kaleido/_kaleido_tab/_tab.py index 5ac7de46..49f0cfee 100644 --- a/src/py/kaleido/_kaleido_tab/_tab.py +++ b/src/py/kaleido/_kaleido_tab/_tab.py @@ -1,10 +1,15 @@ from __future__ import annotations import base64 -from typing import TYPE_CHECKING +import json as _stdlib_json +from typing import TYPE_CHECKING, Any import logistro -import orjson + +try: + import orjson +except ImportError: # pragma: no cover - exercised only when orjson is absent + orjson = None # type: ignore[assignment] from . import _devtools_utils as _dtools from . import _js_logger @@ -32,6 +37,37 @@ def _orjson_default(obj): raise TypeError(f"Type is not JSON serializable: {type(obj).__name__}") +class _StdlibJSONEncoder(_stdlib_json.JSONEncoder): + """ + Encoder used when ``orjson`` is unavailable; mirrors ``_orjson_default``. + + Reproduces the ``orjson.OPT_SERIALIZE_NUMPY`` behavior via the standard + ``.tolist()`` round-trip so callers see the same output regardless of + whether ``orjson`` is installed. + """ + + def default(self, o: Any) -> Any: + if hasattr(o, "tolist"): + return o.tolist() + return super().default(o) + + +def _serialize_spec(spec: Any) -> str: + """ + Serialize a figure spec to a JSON string. + + Uses :mod:`orjson` when available (fast path with native NumPy support); + falls back to the standard-library :mod:`json` module otherwise. + """ + if orjson is not None: + return orjson.dumps( + spec, + default=_orjson_default, + option=orjson.OPT_SERIALIZE_NUMPY, + ).decode() + return _stdlib_json.dumps(spec, cls=_StdlibJSONEncoder) + + def _subscribe_new(tab: choreo.Tab, event: str) -> asyncio.Future: """Create subscription to tab clearing old ones first: helper function.""" new_future = tab.subscribe_once(event) @@ -148,11 +184,7 @@ async def _calc_fig( stepper, ) -> bytes: render_prof.profile_log.tick("serializing spec") - spec_str = orjson.dumps( - spec, - default=_orjson_default, - option=orjson.OPT_SERIALIZE_NUMPY, - ).decode() + spec_str = _serialize_spec(spec) render_prof.profile_log.tick("spec serialized") render_prof.profile_log.tick("sending javascript") diff --git a/src/py/kaleido/mocker/_utils.py b/src/py/kaleido/mocker/_utils.py index d18d9e60..d113f5b1 100644 --- a/src/py/kaleido/mocker/_utils.py +++ b/src/py/kaleido/mocker/_utils.py @@ -1,11 +1,16 @@ from __future__ import annotations import itertools +import json as _stdlib_json from pathlib import Path from typing import TYPE_CHECKING, TypedDict import logistro -import orjson + +try: + import orjson +except ImportError: # pragma: no cover - exercised only when orjson is absent + orjson = None # type: ignore[assignment] from ._args import args @@ -44,30 +49,34 @@ def load_figures_from_paths(paths: list[Path]) -> Generator[FigureDict, None]: if not path.is_file(): raise RuntimeError(f"Path {path} is not a file.") _logger.info(f"Found file: {path!s}") - with path.open(encoding="utf-8") as file: - figure = orjson.loads(file.read()) - for f, w, h, s in itertools.product( # all combos - args.format, - args.width, - args.height, - args.scale, - ): - name = ( - f"{path.stem}.{f!s}" - if not args.parameterize - else f"{path.stem!s}-{w!s}x{h!s}@{s!s}.{f!s}" - ) - opts: LayoutOpts = { - "scale": s, - "width": w, - "height": h, - } - _logger.info(f"Yielding spec: {name!s}") - yield { - "fig": figure, - "path": str(Path(args.output) / name), - "opts": opts, - } + if orjson is not None: + with path.open("rb") as file: + figure = orjson.loads(file.read()) + else: + with path.open(encoding="utf-8") as file: + figure = _stdlib_json.load(file) + for f, w, h, s in itertools.product( # all combos + args.format, + args.width, + args.height, + args.scale, + ): + name = ( + f"{path.stem}.{f!s}" + if not args.parameterize + else f"{path.stem!s}-{w!s}x{h!s}@{s!s}.{f!s}" + ) + opts: LayoutOpts = { + "scale": s, + "width": w, + "height": h, + } + _logger.info(f"Yielding spec: {name!s}") + yield { + "fig": figure, + "path": str(Path(args.output) / name), + "opts": opts, + } class FigureDict(TypedDict): """The type a fig_dicts returns for `write_fig_from_object`.""" diff --git a/src/py/pyproject.toml b/src/py/pyproject.toml index b0532cab..63b4e6a7 100644 --- a/src/py/pyproject.toml +++ b/src/py/pyproject.toml @@ -28,10 +28,12 @@ maintainers = [ dependencies = [ "choreographer>=1.3.0", "logistro>=1.0.8", - "orjson>=3.10.15", "packaging", ] +[project.optional-dependencies] +orjson = ["orjson>=3.10.15"] + [project.urls] Homepage = "https://github.com/plotly/kaleido" Repository = "https://github.com/plotly/kaleido"