WIP: Add decode_timedelta argument to open_virtual_dataset

virtualizarr/manifests/store.py

@@ -3,9 +3,10 @@
 import re
 from collections.abc import AsyncGenerator, Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal, TypeAlias
+from typing import TYPE_CHECKING, Literal, Mapping, TypeAlias
 from urllib.parse import urlparse
 
+from xarray.coders import CFDatetimeCoder, CFTimedeltaCoder
 from zarr.abc.store import (
     ByteRequest,
     OffsetByteRequest,
@@ -281,7 +282,14 @@ def to_virtual_dataset(
         self,
         group="",
         loadable_variables: Iterable[str] | None = None,
-        decode_times: bool | None = None,
+        decode_times: bool
+        | CFDatetimeCoder
+        | Mapping[str, bool | CFDatetimeCoder]
+        | None = None,
+        decode_timedelta: bool
+        | CFTimedeltaCoder
+        | Mapping[str, bool | CFTimedeltaCoder]
+        | None = None,
     ) -> "xr.Dataset":
         """
         Create a "virtual" [xarray.Dataset][] containing the contents of one zarr group.
@@ -294,6 +302,8 @@ def to_virtual_dataset(
         ----------
         group : str
         loadable_variables : Iterable[str], optional
+        decode_times : bool, optional
+        decode_timedelta : bool, optional
 
         Returns
         -------
@@ -312,6 +322,7 @@ def to_virtual_dataset(
             group=group,
             loadable_variables=loadable_variables,
             decode_times=decode_times,
+            decode_timedelta=decode_timedelta,
         )
 
 

virtualizarr/xarray.py

@@ -17,6 +17,7 @@
 import xarray.indexes
 from xarray import DataArray, Dataset, Index, combine_by_coords
 from xarray.backends.common import _find_absolute_paths
+from xarray.coders import CFDatetimeCoder, CFTimedeltaCoder
 from xarray.core import dtypes
 from xarray.core.types import NestedSequence
 from xarray.structure.combine import _infer_concat_order_from_positions, _nested_combine
@@ -41,7 +42,13 @@ def open_virtual_dataset(
     parser: Parser,
     drop_variables: Iterable[str] | None = None,
     loadable_variables: Iterable[str] | None = None,
-    decode_times: bool | None = None,
+    decode_times: bool
+    | CFDatetimeCoder
+    | Mapping[str, bool | CFDatetimeCoder]
+    | None = None,
+    decode_timedelta: bool
+    | CFTimedeltaCoder
+    | Mapping[str, bool | CFTimedeltaCoder] = None,
 ) -> xr.Dataset:
     """
     Open an archival data source as an [xarray.Dataset][] wrapping virtualized zarr arrays.
@@ -76,7 +83,17 @@ def open_virtual_dataset(
         Variables in the data source to load as Dask/NumPy arrays instead of as virtual arrays.
     decode_times
         Bool that is passed into [xarray.open_dataset][]. Allows time to be decoded into a datetime object.
-
+    decode_timedelta
+        Bool, CFTimedeltaCoder, or dict-like, optional
+        If True, decode variables and coordinates with time units in
+        {"days", "hours", "minutes", "seconds", "milliseconds", "microseconds"}
+        into timedelta objects. If False, leave them encoded as numbers.
+        If None (default), assume the same value of ``decode_times``; if
+        ``decode_times`` is a :py:class:`coders.CFDatetimeCoder` instance, this
+        takes the form of a :py:class:`coders.CFTimedeltaCoder` instance with a
+        matching ``time_unit``.
+        Pass a mapping, e.g. ``{"my_variable": False}``,
+        to toggle this feature per-variable individually.
     Returns
     -------
     vds
@@ -93,6 +110,7 @@ def open_virtual_dataset(
     ds = manifest_store.to_virtual_dataset(
         loadable_variables=loadable_variables,
         decode_times=decode_times,
+        decode_timedelta=decode_timedelta,
     )
     return ds.drop_vars(list(drop_variables or ()))
 
@@ -325,6 +343,7 @@ def construct_virtual_dataset(
     group: str | None = None,
     loadable_variables: Iterable[Hashable] | None = None,
     decode_times: bool | None = None,
+    decode_timedelta: bool | None = None,
     reader_options: Optional[dict] = None,
 ) -> xr.Dataset:
     """
@@ -347,6 +366,7 @@ def construct_virtual_dataset(
         zarr_format=3,
         chunks=None,
         decode_times=decode_times,
+        decode_timedelta=decode_timedelta,
     ) as loadable_ds:
         return replace_virtual_with_loadable_vars(
             fully_virtual_ds, loadable_ds, loadable_variables