From ba2b1b5f6d97b22d7129ae4c3743e9a71f253da9 Mon Sep 17 00:00:00 2001
From: Max Jones <14077947+maxrjones@users.noreply.github.com>
Date: Mon, 27 Jan 2025 20:39:37 -0500
Subject: [PATCH] Support Zarr-Python 3

---
 src/titiler/xarray/pyproject.toml       | 15 ++++-
 src/titiler/xarray/titiler/xarray/io.py | 81 +++++--------------------
 2 files changed, 26 insertions(+), 70 deletions(-)

diff --git a/src/titiler/xarray/pyproject.toml b/src/titiler/xarray/pyproject.toml
index e0e1e2112..e7f29cd3b 100644
--- a/src/titiler/xarray/pyproject.toml
+++ b/src/titiler/xarray/pyproject.toml
@@ -38,7 +38,7 @@ dependencies = [
 
 [project.optional-dependencies]
 full = [
-    "zarr>=2,<3",
+    "zarr",
     "h5netcdf",
     "fsspec",
     "s3fs",
@@ -46,7 +46,7 @@ full = [
     "gcsfs",
 ]
 minimal = [
-    "zarr>=2,<3",
+    "zarr",
     "h5netcdf",
     "fsspec",
 ]
@@ -64,7 +64,16 @@ test = [
     "pytest-cov",
     "pytest-asyncio",
     "httpx",
-    "zarr>=2,<3",
+    "zarr",
+    "h5netcdf",
+    "fsspec",
+]
+test-zarr-v2 = [
+    "pytest",
+    "pytest-cov",
+    "pytest-asyncio",
+    "httpx",
+    "zarr<3.0",
     "h5netcdf",
     "fsspec",
 ]
diff --git a/src/titiler/xarray/titiler/xarray/io.py b/src/titiler/xarray/titiler/xarray/io.py
index 370fa5e97..0b619324b 100644
--- a/src/titiler/xarray/titiler/xarray/io.py
+++ b/src/titiler/xarray/titiler/xarray/io.py
@@ -9,6 +9,7 @@
 from morecantile import TileMatrixSet
 from rio_tiler.constants import WEB_MERCATOR_TMS
 from rio_tiler.io.xarray import XarrayReader
+from xarray.namedarray.utils import module_available
 
 
 def xarray_open_dataset(  # noqa: C901
@@ -29,21 +30,6 @@ def xarray_open_dataset(  # noqa: C901
     """
     import fsspec  # noqa
 
-    try:
-        import gcsfs
-    except ImportError:  # pragma: nocover
-        gcsfs = None  # type: ignore
-
-    try:
-        import s3fs
-    except ImportError:  # pragma: nocover
-        s3fs = None  # type: ignore
-
-    try:
-        import aiohttp
-    except ImportError:  # pragma: nocover
-        aiohttp = None  # type: ignore
-
     try:
         import h5netcdf
     except ImportError:  # pragma: nocover
@@ -66,55 +52,8 @@ def xarray_open_dataset(  # noqa: C901
 
     else:
         assert zarr is not None, "'zarr' must be installed to read Zarr dataset"
-
         xr_engine = "zarr"
-
-    if protocol in ["", "file"]:
-        filesystem = fsspec.filesystem(protocol)  # type: ignore
-        file_handler = (
-            filesystem.open(src_path)
-            if xr_engine == "h5netcdf"
-            else filesystem.get_mapper(src_path)
-        )
-
-    elif protocol == "s3":
-        assert (
-            s3fs is not None
-        ), "'aiohttp' must be installed to read dataset stored online"
-
-        s3_filesystem = s3fs.S3FileSystem()
-        file_handler = (
-            s3_filesystem.open(src_path)
-            if xr_engine == "h5netcdf"
-            else s3fs.S3Map(root=src_path, s3=s3_filesystem)
-        )
-
-    elif protocol == "gs":
-        assert (
-            gcsfs is not None
-        ), "'gcsfs' must be installed to read dataset stored in Google Cloud Storage"
-
-        gcs_filesystem = gcsfs.GCSFileSystem()
-        file_handler = (
-            gcs_filesystem.open(src_path)
-            if xr_engine == "h5netcdf"
-            else gcs_filesystem.get_mapper(root=src_path)
-        )
-
-    elif protocol in ["http", "https"]:
-        assert (
-            aiohttp is not None
-        ), "'aiohttp' must be installed to read dataset stored online"
-
-        filesystem = fsspec.filesystem(protocol)  # type: ignore
-        file_handler = (
-            filesystem.open(src_path)
-            if xr_engine == "h5netcdf"
-            else filesystem.get_mapper(src_path)
-        )
-
-    else:
-        raise ValueError(f"Unsupported protocol: {protocol}, for {src_path}")
+        _zarr_v3 = module_available("zarr", minversion="3.0")
 
     # Arguments for xarray.open_dataset
     # Default args
@@ -135,13 +74,21 @@ def xarray_open_dataset(  # noqa: C901
                 "lock": False,
             }
         )
-
-        ds = xarray.open_dataset(file_handler, **xr_open_args)
+        fs = fsspec.filesystem(protocol)
+        ds = xarray.open_dataset(fs.open(src_path), **xr_open_args)
 
     # Fallback to Zarr
     else:
-        ds = xarray.open_zarr(file_handler, **xr_open_args)
-
+        if _zarr_v3:
+            if protocol == "file":
+                store = zarr.storage.LocalStore(parsed.path, read_only=True)
+            else:
+                fs = fsspec.filesystem(protocol, storage_options={"asynchronous": True})
+                store = zarr.storage.FsspecStore(fs, path=src_path, read_only=True)
+            ds = xarray.open_zarr(store, **xr_open_args)
+        else:
+            fs = fsspec.filesystem(protocol)
+            ds = xarray.open_zarr(fs.get_mapper(src_path), **xr_open_args)
     return ds