from __future__ import annotations

import base64

import json

import os

import struct

from collections.abc import Hashable, Iterable, Mapping

from typing import TYPE_CHECKING, Any, Literal, Self, cast, override

import numpy as np

import pandas as pd

from xarray import coding, conventions

from xarray.backends.chunks import grid_rechunk, validate_grid_chunks_alignment

from xarray.backends.common import (

BACKEND_ENTRYPOINTS,

AbstractWritableDataStore,

BackendArray,

BackendEntrypoint,

T_PathFileOrDataStore,

_encode_variable_name,

_normalize_path,

datatree_from_dict_with_io_cleanup,

ensure_dtype_not_object,

)

from xarray.backends.store import StoreBackendEntrypoint

from xarray.conventions import ZARR_CODERS

from xarray.core import indexing

from xarray.core.treenode import NodePath

from xarray.core.types import ZarrWriteModes

from xarray.core.utils import (

Default,

FrozenDict,

HiddenKeyDict,

_default,

attempt_import,

close_on_error,

emit_user_level_warning,

)

from xarray.core.variable import Variable

from xarray.namedarray.parallelcompat import guess_chunkmanager

from xarray.namedarray.pycompat import integer_types

if TYPE_CHECKING:

from zarr import AsyncArray

from xarray.coders import CFDatetimeCoder, CFTimedeltaCoder

from xarray.core.dataset import Dataset

from xarray.core.datatree import DataTree

from xarray.core.types import T_Chunks, ZarrArray, ZarrGroup, ZarrStoreLike

def _get_mappers(*, storage_options, store):

# expand str and path-like arguments

store = _normalize_path(store)

kwargs = {}

if storage_options is None:

mapper = store

else:

if not isinstance(store, str):

raise ValueError(

f"store must be a string to use storage_options. Got {type(store)}"

)

kwargs["storage_options"] = storage_options

mapper = store

return kwargs, mapper

def _choose_default_mode(

*,

mode: ZarrWriteModes | None,

append_dim: Hashable | None,

region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None,

) -> ZarrWriteModes:

if mode is None:

if append_dim is not None:

mode = "a"

elif region is not None:

mode = "r+"

else:

mode = "w-"

if mode not in ["a", "a-"] and append_dim is not None:

raise ValueError("cannot set append_dim unless mode='a' or mode=None")

if mode not in ["a", "a-", "r+"] and region is not None:

raise ValueError(

"cannot set region unless mode='a', mode='a-', mode='r+' or mode=None"

)

if mode not in ["w", "w-", "a", "a-", "r+"]:

raise ValueError(

"The only supported options for mode are 'w', "

f"'w-', 'a', 'a-', and 'r+', but mode={mode!r}"

)

return mode

# need some special secret attributes to tell us the dimensions

DIMENSION_KEY = "_ARRAY_DIMENSIONS"

ZarrFormat = Literal[2, 3]

class FillValueCoder:

"""Handle custom logic to safely encode and decode fill values in Zarr.

Possibly redundant with logic in xarray/coding/variables.py but needs to be

isolated from NetCDF-specific logic.

"""

@classmethod

def encode(

cls, value: int | float | complex | str | bytes, dtype: np.dtype[Any]

) -> Any:

if dtype.kind == "S":

# byte string, this implies that 'value' must also be `bytes` dtype.

if not isinstance(value, bytes):

raise TypeError(

f"Failed to encode fill_value: expected bytes for dtype {dtype}, got {type(value).__name__}"

)

return base64.standard_b64encode(value).decode()

elif dtype.kind == "b":

# boolean

return bool(value)

elif dtype.kind in "iu":

if not isinstance(value, int | float | np.integer | np.floating):

raise TypeError(

f"Failed to encode fill_value: expected int or float for dtype {dtype}, got {type(value).__name__}"

)

return int(value)

elif dtype.kind == "f":

if not isinstance(value, int | float | np.integer | np.floating):

raise TypeError(

f"Failed to encode fill_value: expected int or float for dtype {dtype}, got {type(value).__name__}"

)

return base64.standard_b64encode(struct.pack("<d", float(value))).decode()

elif dtype.kind == "c":

# complex - encode each component as base64, matching float encoding

if not isinstance(value, complex) and not np.issubdtype(

type(value), np.complexfloating

):

raise TypeError(

f"Failed to encode fill_value: expected complex for dtype {dtype}, got {type(value).__name__}"

)

return [

base64.standard_b64encode(

struct.pack("<d", float(value.real)) # type: ignore[union-attr]

).decode(),

base64.standard_b64encode(

struct.pack("<d", float(value.imag)) # type: ignore[union-attr]

).decode(),

]

elif dtype.kind == "U":

return str(value)

else:

raise ValueError(f"Failed to encode fill_value. Unsupported dtype {dtype}")

@classmethod

def decode(

cls, value: int | float | str | bytes | list | None, dtype: str | np.dtype[Any]

):

if value is None:

return None

if dtype == "string":

# zarr V3 string type

return str(value)

elif dtype == "bytes":

# zarr V3 bytes type

if not isinstance(value, str | bytes):

raise TypeError(

f"Failed to decode fill_value: expected str or bytes for dtype {dtype}, got {type(value).__name__}"

)

return base64.standard_b64decode(value)

np_dtype = np.dtype(dtype)

if (

np_dtype.kind == "S" or coding.strings.is_unicode_dtype(np_dtype)

) and isinstance(value, str):

return np.asarray(value, dtype=np_dtype)[()]

if np_dtype.kind == "f":

if isinstance(value, int | float) and not isinstance(value, bool):

return np.asarray(value, dtype=np_dtype)[()]

if not isinstance(value, str | bytes):

raise TypeError(

f"Failed to decode fill_value: expected str or bytes for dtype {np_dtype}, got {type(value).__name__}"

)

return struct.unpack("<d", base64.standard_b64decode(value))[0]

elif np_dtype.kind == "c":

# complex - decode each component from base64, matching float decoding

if not (isinstance(value, list | tuple) and len(value) == 2):

raise TypeError(

f"Failed to decode fill_value: expected a 2-element list for dtype {np_dtype}, got {type(value).__name__}"

)

real = struct.unpack("<d", base64.standard_b64decode(value[0]))[0]

imag = struct.unpack("<d", base64.standard_b64decode(value[1]))[0]

return complex(real, imag)

elif np_dtype.kind == "b":

return bool(value)

elif np_dtype.kind in "iu":

if not isinstance(value, int | float | np.integer | np.floating):

raise TypeError(

f"Failed to decode fill_value: expected int or float for dtype {np_dtype}, got {type(value).__name__}"

)

return int(value)

else:

raise ValueError(f"Failed to decode fill_value. Unsupported dtype {dtype}")

def encode_zarr_attr_value(value):

"""Encode an attribute value as something that can be serialized as json.

Many xarray datasets / variables have numpy arrays and values. This

function handles encoding / decoding of such items.

ndarray -> list

scalar array -> scalar

other -> other (no change)

"""

if isinstance(value, np.ndarray):

encoded = value.tolist()

elif isinstance(value, np.generic):

encoded = value.item()

else:

encoded = value

return encoded

def has_zarr_async_index() -> bool:

try:

import zarr

return hasattr(zarr.AsyncArray, "oindex")

except (ImportError, AttributeError):

return False

class ZarrArrayWrapper(BackendArray):

__slots__ = ("_array", "dtype", "shape")

def __init__(self, zarr_array):

# some callers attempt to evaluate an array if an `array` property exists on the object.

# we prefix with _ to avoid this inference.

# TODO type hint this?

self._array = zarr_array

self.shape = self._array.shape

# preserve vlen string object dtype (GH 7328), but keep a native StringDType

dtype = self._array.dtype

if (

dtype.kind != "T"

and self._array.serializer

and self._array.serializer.to_dict()["name"] == "vlen-utf8"

):

dtype = coding.strings.create_vlen_dtype(str)

self.dtype = dtype

def get_array(self):

return self._array

def _preserve_string_dtype(self, value):

if self.dtype.kind == "T" and isinstance(value, str):

return np.asarray(value, dtype=self.dtype)

return value

def _oindex(self, key):

return self._array.oindex[key]

def _vindex(self, key):

return self._array.vindex[key]

def _getitem(self, key):

return self._array[key]

@property

def _async_array(self) -> AsyncArray[Any]:

# public since zarr 3.1.4

if hasattr(self._array, "async_array"):

return self._array.async_array

return self._array._async_array

async def _async_getitem(self, key):

async_array = self._async_array

return await async_array.getitem(key)

async def _async_oindex(self, key):

if not has_zarr_async_index():

raise NotImplementedError(

"For lazy orthogonal async indexing with zarr, zarr-python=>v3.1.2 is required"

)

async_array = self._async_array

return await async_array.oindex.getitem(key)

async def _async_vindex(self, key):

if not has_zarr_async_index():

raise NotImplementedError(

"For lazy vectorized async indexing with zarr, zarr-python=>v3.1.2 is required"

)

async_array = self._async_array

return await async_array.vindex.getitem(key)

def __getitem__(self, key):

array = self._array

if isinstance(key, indexing.BasicIndexer):

method = self._getitem

elif isinstance(key, indexing.VectorizedIndexer):

method = self._vindex

elif isinstance(key, indexing.OuterIndexer):

method = self._oindex

value = indexing.explicit_indexing_adapter(

key, array.shape, indexing.IndexingSupport.VECTORIZED, method

)

# if self.ndim == 0:

# could possibly have a work-around for 0d data here

return self._preserve_string_dtype(value)

@override

async def async_getitem(self, key):

array = self._array

if isinstance(key, indexing.BasicIndexer):

method = self._async_getitem

elif isinstance(key, indexing.VectorizedIndexer):

method = self._async_vindex

elif isinstance(key, indexing.OuterIndexer):

method = self._async_oindex

value = await indexing.async_explicit_indexing_adapter(

key, array.shape, indexing.IndexingSupport.VECTORIZED, method

)

return self._preserve_string_dtype(value)

def _rectilinear_encoding_error(key: str, value, name) -> TypeError:

"""Error for a rectilinear (variable-sized) chunk or shard spec.

xarray can read such specs but not yet write them.

"""

return TypeError(

f"encoding[{key!r}]={value!r} for variable {name!r} is a rectilinear "

"(variable-sized) grid, e.g. read from a store opened with "

"zarr.config.set({'array.rectilinear_chunks': True}). xarray can read "

"rectilinear grids but not yet write them, including with `region` or "

"`append_dim`. To write a regular grid to a new array instead, clear "

f"the encoding (`del ds[{name!r}].encoding[{key!r}]`) and make sure the "

"variable's own chunks are uniform (e.g. `ds.chunk({dim: size})`)."

)

def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name):

"""Determine the zarr chunks for a variable.

Given encoding chunks (possibly None or []) and variable chunks

(possibly None or []).

"""

# zarr chunk spec:

# chunks : int or tuple of ints, optional

# Chunk shape. If not provided, will be guessed from shape and dtype.

# if there are no chunks in encoding and the variable data is a numpy

# array, then we let zarr use its own heuristics to pick the chunks

if not var_chunks and not enc_chunks:

return None

# if there are no chunks in encoding but there are dask chunks, we try to

# use the same chunks in zarr

# However, zarr chunks needs to be uniform for each array

# https://zarr-specs.readthedocs.io/en/latest/v2/v2.0.html#chunks

# while dask chunks can be variable sized

# https://dask.pydata.org/en/latest/array-design.html#chunks

if var_chunks and not enc_chunks:

if any(len(set(chunks[:-1])) > 1 for chunks in var_chunks):

raise ValueError(

"Zarr requires uniform chunk sizes except for final chunk. "

f"Variable named {name!r} has incompatible dask chunks: {var_chunks!r}. "

"Consider rechunking using `chunk()`."

)

if any((chunks[0] < chunks[-1]) for chunks in var_chunks):

raise ValueError(

"Final chunk of Zarr array must be the same size or smaller "

f"than the first. Variable named {name!r} has incompatible Dask chunks {var_chunks!r}."

"Consider either rechunking using `chunk()` or instead deleting "

"or modifying `encoding['chunks']`."

)

# return the first chunk for each dimension

return tuple(chunk[0] for chunk in var_chunks)

# From here on, we are dealing with user-specified chunks in encoding

# zarr allows chunks to be an integer, in which case it uses the same chunk

# size on each dimension.

# Here we re-implement this expansion ourselves. That makes the logic of

# checking chunk compatibility easier

if isinstance(enc_chunks, integer_types):

enc_chunks_tuple = ndim * (enc_chunks,)

else:

enc_chunks_tuple = tuple(enc_chunks)

if len(enc_chunks_tuple) != ndim:

# throw away encoding chunks, start over

return _determine_zarr_chunks(

None,

var_chunks,

ndim,

name,

)

for x in enc_chunks_tuple:

if not isinstance(x, int):

if isinstance(x, list | tuple):

raise _rectilinear_encoding_error("chunks", enc_chunks_tuple, name)

raise TypeError(

"zarr chunk sizes specified in `encoding['chunks']` "

"must be an int or a tuple of ints. "

f"Instead found encoding['chunks']={enc_chunks_tuple!r} "

f"for variable named {name!r}."

)

# if there are chunks in encoding and the variable data is a numpy array,

# we use the specified chunks

if not var_chunks:

return enc_chunks_tuple

return enc_chunks_tuple

def _compact_chunk_sizes(

chunk_sizes: tuple[tuple[int, ...], ...],

) -> tuple[int | tuple[int, ...], ...]:

"""Simplify a dask-style listing of chunk sizes along regular dimensions.

Replace a dask-style listing of chunk sizes with a single int along any

dimension where it describes a regular grid, e.g. ((10, 10, 5),) -> (10,).

Gives a consistent representation across zarr-python versions.

"""

compacted: list[int | tuple[int, ...]] = []

for sizes in chunk_sizes:

sizes = tuple(sizes)

if not sizes:

# zero-length dimension: dask rejects an empty tuple here

compacted.append((0,))

elif len(set(sizes[:-1])) <= 1 and sizes[-1] <= sizes[0]:

compacted.append(sizes[0])

else:

compacted.append(sizes)

return tuple(compacted)

def _get_zarr_dims_and_attrs(zarr_obj, dimension_key, try_nczarr):

# Check for attributes and dimension name metadata as discussed in the Zarr encoding

# specification https://docs.xarray.dev/en/stable/internals/zarr-encoding-spec.html

# Zarr V3 explicitly stores the dimension names in the metadata

try:

# if this exists, we are looking at a Zarr V3 array

# convert None to empty tuple

dimensions = zarr_obj.metadata.dimension_names or ()

except AttributeError:

# continue to old code path

pass

else:

attributes = dict(zarr_obj.attrs)

if len(zarr_obj.shape) != len(dimensions):

raise KeyError(

"Zarr object is missing the `dimension_names` metadata which is "

"required for xarray to determine variable dimensions."

)

return dimensions, attributes

# Zarr arrays do not have dimensions. To get around this problem, we add

# an attribute that specifies the dimension. We have to hide this attribute

# when we send the attributes to the user.

# zarr_obj can be either a zarr group or zarr array

try:

# Xarray-Zarr

dimensions = zarr_obj.attrs[dimension_key]

except KeyError as e:

if not try_nczarr:

raise KeyError(

f"Zarr object is missing the attribute `{dimension_key}`, which is "

"required for xarray to determine variable dimensions."

) from e

# NCZarr defines dimensions through metadata in .zarray

zarray_path = os.path.join(zarr_obj.path, ".zarray")

import asyncio

zarray_str = asyncio.run(zarr_obj.store.get(zarray_path)).to_bytes()

zarray = json.loads(zarray_str)

try:

# NCZarr uses Fully Qualified Names

dimensions = [

os.path.basename(dim) for dim in zarray["_NCZARR_ARRAY"]["dimrefs"]

]

except KeyError as e:

raise KeyError(

f"Zarr object is missing the attribute `{dimension_key}` and the NCZarr metadata, "

"which are required for xarray to determine variable dimensions."

) from e

nc_attrs = [attr for attr in zarr_obj.attrs if attr.lower().startswith("_nc")]

attributes = HiddenKeyDict(zarr_obj.attrs, [dimension_key] + nc_attrs)

return dimensions, attributes

def extract_zarr_variable_encoding(

variable,

raise_on_invalid=False,

name=None,

*,

zarr_format: ZarrFormat,

):

"""Extract zarr encoding dictionary from xarray Variable.

Parameters

----------

variable : Variable

raise_on_invalid : bool, optional

name: str | Hashable, optional

zarr_format: Literal[2,3]

Returns

-------

encoding : dict

Zarr encoding for `variable`

"""

encoding = variable.encoding.copy()

safe_to_drop = {"source", "original_shape", "preferred_chunks"}

valid_encodings = {

"chunks",

"shards",

"compressors",

"filters",

"serializer",

"cache_metadata",

"write_empty_chunks",

"chunk_key_encoding",

}

if zarr_format == 3:

valid_encodings.add("fill_value")

for k in safe_to_drop:

if k in encoding:

del encoding[k]

if raise_on_invalid:

invalid = [k for k in encoding if k not in valid_encodings]

if "fill_value" in invalid and zarr_format == 2:

msg = " Use `_FillValue` to set the Zarr array `fill_value`"

else:

msg = ""

if invalid:

raise ValueError(

f"unexpected encoding parameters for zarr backend: {invalid!r}." + msg

)

else:

for k in list(encoding):

if k not in valid_encodings:

del encoding[k]

chunks = _determine_zarr_chunks(

enc_chunks=encoding.get("chunks"),

var_chunks=variable.chunks,

ndim=variable.ndim,

name=name,

)

if chunks is None:

chunks = "auto"

encoding["chunks"] = chunks

# Reject rectilinear shards (a sequence of sequences) like we do for chunks.

# Zarr also accepts an int, tuple of ints, "auto" or a dict here.

shards = encoding.get("shards")

if isinstance(shards, list | tuple) and any(

isinstance(x, list | tuple) for x in shards

):

raise _rectilinear_encoding_error("shards", shards, name)

if isinstance(shards, integer_types):

# Expand to a tuple: zarr-python 3.2.x crashes on an int shard spec.

encoding["shards"] = variable.ndim * (int(shards),)

return encoding

# Function below is copied from conventions.encode_cf_variable.

# The only change is to raise an error for object dtypes.

def encode_zarr_variable(

var, needs_copy=True, name=None, *, zarr_format: ZarrFormat | None = None

):

"""Convert a Variable into another Variable which follows the CF conventions.

In particular:

- Nans are masked using _FillValue (or the deprecated missing_value)

- Rescaling via: scale_factor and add_offset

- datetimes are converted to the CF 'units since time' format

- dtype encodings are enforced.

Parameters

----------

var : Variable

A variable holding un-encoded data.

Returns

-------

out : Variable

A variable which has been encoded as described above.

"""

var = conventions.encode_cf_variable(var, name=name, coders=ZARR_CODERS)

var = ensure_dtype_not_object(var, name=name)

if (

zarr_format == 3

and var.dtype.kind == "T"

and var.dtype == np.dtypes.StringDType()

):

return var

# zarr allows unicode, but not variable-length strings, so it's both

# simpler and more compact to always encode as UTF-8 explicitly.

# TODO: allow toggling this explicitly via dtype in encoding.

# TODO: revisit this now that Zarr _does_ allow variable-length strings

coder = coding.strings.EncodedStringCoder(allows_unicode=True)

var = coder.encode(var, name=name)

var = coding.strings.ensure_fixed_length_bytes(var)

return var

def _validate_datatypes_for_zarr_append(vname, existing_var, new_var):

"""Check that the dtype of the data to append is compatible with the store.

If variable exists in the store, confirm dtype of the data to append is

compatible with existing dtype.

"""

if (

np.issubdtype(new_var.dtype, np.number)

or np.issubdtype(new_var.dtype, np.datetime64)

or np.issubdtype(new_var.dtype, np.bool_)

or new_var.dtype == object

or (new_var.dtype.kind in ("S", "U") and existing_var.dtype == object)

):

# We can skip dtype equality checks under two conditions: (1) if the var to append is

# new to the dataset, because in this case there is no existing var to compare it to;

# or (2) if var to append's dtype is known to be easy-to-append, because in this case

# we can be confident appending won't cause problems. Examples of dtypes which are not

# easy-to-append include length-specified strings of type `|S*` or `<U*` (where * is a

# positive integer character length). For these dtypes, appending dissimilar lengths

# can result in truncation of appended data. Therefore, variables which already exist

# in the dataset, and with dtypes which are not known to be easy-to-append, necessitate

# exact dtype equality, as checked below.

pass

elif new_var.dtype != existing_var.dtype:

raise ValueError(

f"Mismatched dtypes for variable {vname} between Zarr store on disk "

f"and dataset to append. Store has dtype {existing_var.dtype} but "

f"dataset to append has dtype {new_var.dtype}."

)

def _validate_and_transpose_existing_dims(

var_name, new_var, existing_var, region, append_dim

):

if new_var.dims != existing_var.dims:

if set(existing_var.dims) == set(new_var.dims):

new_var = new_var.transpose(*existing_var.dims)

else:

raise ValueError(

f"variable {var_name!r} already exists with different "

f"dimension names {existing_var.dims} != "

f"{new_var.dims}, but changing variable "

f"dimensions is not supported by to_zarr()."

)

existing_sizes = {}

for dim, size in existing_var.sizes.items():

if region is not None and dim in region:

start, stop, stride = region[dim].indices(size)

assert stride == 1 # region was already validated

size = stop - start

if dim != append_dim:

existing_sizes[dim] = size

new_sizes = {dim: size for dim, size in new_var.sizes.items() if dim != append_dim}

if existing_sizes != new_sizes:

raise ValueError(

f"variable {var_name!r} already exists with different "

f"dimension sizes: {existing_sizes} != {new_sizes}. "

f"to_zarr() only supports changing dimension sizes when "

f"explicitly appending, but append_dim={append_dim!r}. "

f"If you are attempting to write to a subset of the "

f"existing store without changing dimension sizes, "

f"consider using the region argument in to_zarr()."

)

return new_var

def _put_attrs(zarr_obj, attrs):

"""Raise a more informative error message for invalid attrs."""

try:

zarr_obj.attrs.put(attrs)

except TypeError as e:

raise TypeError("Invalid attribute in Dataset.attrs.") from e

return zarr_obj

class ZarrStore(AbstractWritableDataStore):

"""Store for reading and writing data via zarr."""

__slots__ = (

"_align_chunks",

"_append_dim",

"_cache_members",

"_close_store_on_close",

"_consolidate_on_close",

"_group",

"_members",

"_mode",

"_read_only",

"_safe_chunks",

"_use_zarr_fill_value_as_mask",

"_write_empty",

"_write_region",

"zarr_group",

)

@classmethod

def open_store(

cls,

store,

mode: ZarrWriteModes = "r",

group=None,

consolidated=False,

consolidate_on_close=False,

storage_options=None,

append_dim=None,

write_region=None,

safe_chunks=True,

align_chunks=False,

zarr_format=None,

use_zarr_fill_value_as_mask=None,

write_empty: bool | None = None,

cache_members: bool = True,

):

(

zarr_group,

consolidate_on_close,

close_store_on_close,

use_zarr_fill_value_as_mask,

) = _get_open_params(

store=store,

mode=mode,

group=group,

consolidated=consolidated,

consolidate_on_close=consolidate_on_close,

storage_options=storage_options,

use_zarr_fill_value_as_mask=use_zarr_fill_value_as_mask,

zarr_format=zarr_format,

)

from zarr import Group

group_members: dict[str, Group] = {}

group_paths = list(_iter_zarr_groups(zarr_group, parent=group))

for path in group_paths:

if path == group:

group_members[path] = zarr_group

else:

rel_path = path.removeprefix(f"{group}/")

group_members[path] = zarr_group[rel_path.removeprefix("/")]

out = {

group: cls(

group_store,

mode,

consolidate_on_close,

append_dim,

write_region,

safe_chunks,

write_empty,

close_store_on_close,

use_zarr_fill_value_as_mask,

align_chunks=align_chunks,

cache_members=cache_members,

)

for group, group_store in group_members.items()

}

return out

@classmethod

def open_group(

cls,

store,

mode: ZarrWriteModes = "r",

group=None,

consolidated=False,

consolidate_on_close=False,

storage_options=None,

append_dim=None,

write_region=None,

safe_chunks=True,

align_chunks=False,

zarr_format=None,

use_zarr_fill_value_as_mask=None,

write_empty: bool | None = None,

cache_members: bool = True,

):

(

zarr_group,

consolidate_on_close,

close_store_on_close,

use_zarr_fill_value_as_mask,

) = _get_open_params(

store=store,

mode=mode,

group=group,

consolidated=consolidated,

consolidate_on_close=consolidate_on_close,

storage_options=storage_options,

use_zarr_fill_value_as_mask=use_zarr_fill_value_as_mask,

zarr_format=zarr_format,

)

return cls(

zarr_group,

mode,

consolidate_on_close,

append_dim,

write_region,

safe_chunks,

write_empty,

close_store_on_close,

use_zarr_fill_value_as_mask,

align_chunks=align_chunks,

cache_members=cache_members,

)

def __init__(

self,

zarr_group,

mode=None,

consolidate_on_close=False,

append_dim=None,

write_region=None,

safe_chunks=True,

write_empty: bool | None = None,

close_store_on_close: bool = False,

use_zarr_fill_value_as_mask=None,

align_chunks: bool = False,

cache_members: bool = True,

):

if align_chunks:

# Disabled the safe_chunks validations if the alignment is going to be applied

safe_chunks = False

self.zarr_group = zarr_group

self._read_only = self.zarr_group.read_only

self._group = self.zarr_group.path

self._mode = mode

self._consolidate_on_close = consolidate_on_close

self._append_dim = append_dim

self._write_region = write_region

self._align_chunks = align_chunks

self._safe_chunks = safe_chunks

self._write_empty = write_empty

self._close_store_on_close = close_store_on_close

self._use_zarr_fill_value_as_mask = use_zarr_fill_value_as_mask

self._cache_members: bool = cache_members

self._members: dict[str, ZarrArray | ZarrGroup] = {}

if self._cache_members:

# initialize the cache

# this cache is created here and never updated.

# If the `ZarrStore` instance creates a new zarr array, or if an external process

# removes an existing zarr array, then the cache will be invalid.

# We use this cache only to record any pre-existing arrays when the group was opened

# create a new ZarrStore instance if you want to

# capture the current state of the zarr group, or create a ZarrStore with

# `cache_members` set to `False` to disable this cache and instead fetch members

# on demand.

self._members = self._fetch_members()

@override

def get_child_store(self, group: str) -> Self:

zarr_group = self.zarr_group.require_group(group)

return type(self)(

zarr_group=zarr_group,

mode=self._mode,

consolidate_on_close=self._consolidate_on_close,

append_dim=self._append_dim,

write_region=self._write_region,

safe_chunks=self._safe_chunks,

write_empty=self._write_empty,

close_store_on_close=self._close_store_on_close,

use_zarr_fill_value_as_mask=self._use_zarr_fill_value_as_mask,

align_chunks=self._align_chunks,

cache_members=self._cache_members,

)

@property

def members(self) -> dict[str, ZarrArray | ZarrGroup]:

"""Model the arrays and groups contained in self.zarr_group as a dict.

If `self._cache_members` is true, the dict is cached. Otherwise, it is

retrieved from storage.

"""

if not self._cache_members:

return self._fetch_members()

else:

return self._members

def _fetch_members(self) -> dict[str, ZarrArray | ZarrGroup]:

"""Get the arrays and groups defined in the zarr group of this Store."""

return dict(self.zarr_group.members())

def array_keys(self) -> tuple[str, ...]:

from zarr import Array as ZarrArray

return tuple(

key for (key, node) in self.members.items() if isinstance(node, ZarrArray)

)

def arrays(self) -> tuple[tuple[str, ZarrArray], ...]:

from zarr import Array as ZarrArray

return tuple(

(key, node)

for (key, node) in self.members.items()

if isinstance(node, ZarrArray)

)

@property

def ds(self):

# TODO: consider deprecating this in favor of zarr_group

return self.zarr_group

def open_store_variable(self, name):

zarr_array = self.members[name]

try_nczarr = self._mode == "r"

dimensions, attributes = _get_zarr_dims_and_attrs(

zarr_array, DIMENSION_KEY, try_nczarr

)

attributes = dict(attributes)

array_wrapper = ZarrArrayWrapper(zarr_array)

data = indexing.LazilyIndexedArray(array_wrapper)

try:

chunks = tuple(zarr_array.chunks)

except NotImplementedError:

# Rectilinear chunk grid (zarr-python >= 3.2): `.chunks` raises, so

# read the per-chunk sizes instead, e.g. ((10, 20, 30),).

chunks = zarr_array.read_chunk_sizes

# Normalise to an int per regular dim and a tuple per rectilinear dim,

# since what zarr-python returns above varies between versions.

chunks = _compact_chunk_sizes(

tuple(x if isinstance(x, tuple) else (x,) for x in chunks)

)

try:

shards = zarr_array.shards

except NotImplementedError:

# Rectilinear shard grid: `.shards` raises, so read the per-shard

# (i.e. outer/storage chunk) sizes instead, e.g. ((1, 2),).

shards = _compact_chunk_sizes(zarr_array.write_chunk_sizes)

outer_chunks = shards if shards is not None else chunks

preferred_chunks = dict(zip(dimensions, outer_chunks, strict=True))

encoding = {

"chunks": chunks,

"preferred_chunks": preferred_chunks,

"compressors": zarr_array.compressors,

"filters": zarr_array.filters,

"shards": shards,

}

if array_wrapper.dtype.kind == "T":

encoding["dtype"] = array_wrapper.dtype

if self.zarr_group.metadata.zarr_format == 3:

encoding.update({"serializer": zarr_array.serializer})

if self._use_zarr_fill_value_as_mask:

# Setting this attribute triggers CF decoding for missing values

# by interpreting Zarr's fill_value to mean the same as netCDF's _FillValue

if zarr_array.fill_value is not None:

attributes["_FillValue"] = zarr_array.fill_value

else: