Fixes retrieval of chunk sizes for datasets (#83)

* fixed retrieval of chunk sizes for datasets

* Update core.py

* updated wgats-new.rst

docs/whats-new.rst

@@ -1,6 +1,15 @@
 What's New
 ==========
-v0.3.0 (unreleased)
+v0.3.1
+-------------------
+
+Bug fixes
+~~~~~~~~~~~~~
+- Fixed a bug that prevented parallel frame saving to work with `xarray.Datasets` (:pull:`83`).
+  By `Tomas Chor <https://github.com/tomchor>`.
+
+
+v0.3.0 (2022/3/17)
 -------------------
 
 Documentation

xmovie/core.py

@@ -304,9 +304,7 @@ def __init__(
         if input_check:
             if isinstance(self.data, xr.Dataset):
                 raise ValueError(
-                    "xmovie presets do not yet support the input of xr.Datasets. \
-                In order to use datasets as inputs, set `input_check` to False. \
-                Note that this requires you to manually set colorlimits etc."
+                    "xmovie presets do not yet fully support the input of xr.Datasets.\nIn order to use datasets as inputs, set `input_check` to False.\nNote that this requires you to manually set colorlimits etc."
                 )
 
             # Set defaults
@@ -414,7 +412,13 @@ def save_frames_parallel(self, odir, parallel_compute_kwargs=dict()):
             raise ValueError(
                 f"Input data needs to be a dask array to save in parallel. Please chunk the input with single chunks along {framedim}."
             )
-        framedim_chunks = da.chunks[da.dims.index(framedim)]
+
+        if type(da) is xr.DataArray:
+            framedim_chunks = da.chunks[da.dims.index(framedim)]
+        elif type(da) is xr.Dataset:
+            framedim_chunks = da.chunks[framedim]
+        else:
+            raise(TypeError("`da` must be either an xarray.DataArray or xarray.Dataset"))
 
         if not all([chunk == 1 for chunk in framedim_chunks]):
             raise ValueError(