Do not raise if custom renderer

ahuang11 · Aug 2, 2024 · ffdbada · ffdbada
1 parent 92d2c5e
commit ffdbada
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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")