From 2af50080a888aa119cf8db5edf54e9472a35cd5c Mon Sep 17 00:00:00 2001
From: Simon Perkins <simon.perkins@gmail.com>
Date: Wed, 11 Sep 2024 09:08:45 +0200
Subject: [PATCH] Also add to testing group

---
 pyproject.toml    |  2 +-
 xarray_ms/core.py | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 968f8d9..2e65841 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,7 @@ typing-extensions = { version = "^4.12.2", python = "<3.11" }
 zarr = {version = "^2.18.3", optional = true, extras = ["testing"]}
 
 [tool.poetry.extras]
-testing = ["dask", "distributed", "pytest"]
+testing = ["dask", "distributed", "pytest", "zarr"]
 
 [tool.poetry.plugins."xarray.backends"]
 "xarray-ms:msv2" = "xarray_ms.backend.msv2.entrypoint:MSv2PartitionEntryPoint"
diff --git a/xarray_ms/core.py b/xarray_ms/core.py
index c986b08..b1418cd 100644
--- a/xarray_ms/core.py
+++ b/xarray_ms/core.py
@@ -5,6 +5,11 @@
 from xarray.core.dataset import Dataset
 from xarray.core.datatree import DataTree
 
+try:
+  import zarr
+except ImportError:
+  zarr = None
+
 
 def encode_attributes(ds: Dataset) -> Dataset:
   """Encode the antenna_xds attribute of a Dataset."""
@@ -44,6 +49,9 @@ def xds_from_zarr(*args, **kwargs):
   """Read a Measurement Set-like :class:`~xarray.Dataset` from a Zarr store.
 
   Thin wrapper around :func:`xarray.open_zarr`."""
+  if zarr is None:
+    raise ImportError("pip install zarr")
+
   return decode_attributes(xarray.open_zarr(*args, **kwargs))
 
 
@@ -52,6 +60,9 @@ def xds_to_zarr(ds: Dataset, *args, **kwargs) -> None:
 
   Thin wrapper around :meth:`xarray.Dataset.to_zarr`.
   """
+  if zarr is None:
+    raise ImportError("pip install zarr")
+
   return encode_attributes(ds).to_zarr(*args, **kwargs)
 
 
@@ -60,6 +71,9 @@ def xdt_from_zarr(*args, **kwargs):
   from a Zarr store.
 
   Thin wrapper around :func:`xarray.backends.api.open_datatree`."""
+  if zarr is None:
+    raise ImportError("pip install zarr")
+
   return open_datatree(*args, **kwargs).map_over_subtree(decode_attributes)
 
 
@@ -69,4 +83,7 @@ def xdt_to_zarr(dt: DataTree, *args, **kwargs) -> None:
 
   Thin wrapper around :meth:`xarray.core.datatree.DataTree.to_zarr`.
   """
+  if zarr is None:
+    raise ImportError("pip install zarr")
+
   return dt.map_over_subtree(encode_attributes).to_zarr(*args, **kwargs)