From ffdbada98eb6c2ba89f71aabcb4559a664687bd9 Mon Sep 17 00:00:00 2001 From: ahuang11 Date: Thu, 1 Aug 2024 17:49:34 -0700 Subject: [PATCH] Do not raise if custom renderer --- streamjoy/_utils.py | 3 ++- streamjoy/serializers.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/streamjoy/_utils.py b/streamjoy/_utils.py index 3cb5900..0699b13 100644 --- a/streamjoy/_utils.py +++ b/streamjoy/_utils.py @@ -249,6 +249,7 @@ def validate_xarray( dim: str | None = None, var: str | None = None, warn: bool = True, + raise_ndim: bool = True, ): import xarray as xr @@ -262,7 +263,7 @@ def validate_xarray( squeeze_dims = [d for d in ds.dims if d != dim and ds.sizes[d] == 1] ds = ds.squeeze(squeeze_dims) - if ds.ndim > 3: + if ds.ndim > 3 and raise_ndim: raise ValueError(f"Can only handle 3D arrays; {ds.ndim}D array found") return ds diff --git a/streamjoy/serializers.py b/streamjoy/serializers.py index 268e3b7..cf4cae2 100644 --- a/streamjoy/serializers.py +++ b/streamjoy/serializers.py @@ -119,7 +119,7 @@ def serialize_xarray( dim = kwargs.pop("dim", None) var = kwargs.pop("var", None) - ds = _utils.validate_xarray(ds, dim=dim, var=var) + ds = _utils.validate_xarray(ds, dim=dim, var=var, raise_ndim=renderer is None) if not dim: dim = list(ds.dims)[0] _utils.warn_default_used("dim", dim, suffix="from the dataset")