From 975dada65ee320f0b90dad5b825267b4c523e727 Mon Sep 17 00:00:00 2001
From: Kirill Kouzoubov <kirill888@gmail.com>
Date: Mon, 24 Jun 2024 14:55:31 +1000
Subject: [PATCH] refactor: in zarr driver

- some names not longer fit
- include reference to current driver in the
  loader context
---
 odc/loader/_zarr.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/odc/loader/_zarr.py b/odc/loader/_zarr.py
index 10d7a81..53639a6 100644
--- a/odc/loader/_zarr.py
+++ b/odc/loader/_zarr.py
@@ -98,14 +98,14 @@ class XrMDPlugin:
     def __init__(
         self,
         template: RasterGroupMetadata,
-        src: xr.Dataset | None = None,
+        fallback: xr.Dataset | None = None,
     ) -> None:
         self._template = template
-        self._src = src
+        self._fallback = fallback
 
     def _resolve_src(self, md: Any, regen_coords: bool = False) -> xr.Dataset | None:
         return _resolve_src_dataset(
-            md, regen_coords=regen_coords, fallback=self._src, chunks={}
+            md, regen_coords=regen_coords, fallback=self._fallback, chunks={}
         )
 
     def extract(self, md: Any) -> RasterGroupMetadata:
@@ -150,9 +150,11 @@ def __init__(
         self,
         geobox: GeoBox,
         chunks: None | dict[str, int],
+        driver: Any | None = None,
     ) -> None:
         self.geobox = geobox
         self.chunks = chunks
+        self.driver = driver
 
     def with_env(self, env: dict[str, Any]) -> "Context":
         assert isinstance(env, dict)
@@ -161,7 +163,10 @@ def with_env(self, env: dict[str, Any]) -> "Context":
 
 class XrMemReader:
     """
-    Protocol for raster readers.
+    Implements protocol for raster readers.
+
+    - Read from in-memory xarray.Dataset
+    - Read from zarr spec
     """
 
     # pylint: disable=too-few-public-methods
@@ -242,7 +247,7 @@ def read(
 
 class XrMemReaderDriver:
     """
-    Read from in memory xarray.Dataset.
+    Read from in memory xarray.Dataset or zarr spec document.
     """
 
     Reader = XrMemReader
@@ -265,7 +270,7 @@ def new_load(
         *,
         chunks: None | dict[str, int] = None,
     ) -> Context:
-        return Context(geobox, chunks)
+        return Context(geobox, chunks, driver=self)
 
     def finalise_load(self, load_state: Context) -> Context:
         return load_state
@@ -284,7 +289,7 @@ def open(self, src: RasterSource, ctx: Context) -> XrMemReader:
 
     @property
     def md_parser(self) -> MDParser:
-        return XrMDPlugin(self.template, src=self.src)
+        return XrMDPlugin(self.template, fallback=self.src)
 
     @property
     def dask_reader(self) -> DaskRasterReader | None: