tonik 0.0.11__tar.gz → 0.0.12__tar.gz

{tonik-0.0.11 → tonik-0.0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: tonik
-Version: 0.0.11
+Version: 0.0.12
 Summary: Store time series data as HDF5 files and access them through an API.
 Project-URL: Homepage, https://tsc-tools.github.io/tonik
 Project-URL: Issues, https://github.com/tsc-tools/tonik/issues

{tonik-0.0.11 → tonik-0.0.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tonik"
-version = "0.0.11"
+version = "0.0.12"
 authors = [
   { name="Yannik Behr", email="y.behr@gns.cri.nz" },
   { name="Christof Mueller", email="c.mueller@gns.cri.nz" }

{tonik-0.0.11 → tonik-0.0.12}/src/tonik/xarray2zarr.py

@@ -6,7 +6,47 @@ import xarray as xr
 logger = logging.getLogger(__name__)
 
 
-def xarray2zarr(xds, path, mode='a'):
+def merge_arrays(xds_old: xr.DataArray, xds_new: xr.DataArray) -> xr.DataArray:
+    """
+    Merge two xarray datasets with the same datetime index.
+
+    Parameters
+    ----------
+    xds_old : xr.DataArray
+        Old array.
+    xds_new : xr.DataArray
+        New array.
+
+    Returns
+    -------
+    xr.DataArray
+        Merged array.
+    """
+    xda_old = xds_old.drop_duplicates(
+        'datetime', keep='last')
+    xda_new = xds_new.drop_duplicates(
+        'datetime', keep='last')
+    xda_new = xda_new.combine_first(xda_old)
+    return xda_new
+
+
+def xarray2zarr(xds: xr.Dataset, path: str, mode: str = 'a'):
+    """
+    Write xarray dataset to zarr files.
+
+    Parameters
+    ----------
+    xds : xr.Dataset
+        Dataset to write.
+    path : str
+        Path to write the dataset.
+    mode : str, optional
+        Write mode, by default 'a'.
+
+    Returns
+    -------
+    None
+    """
     for feature in xds.data_vars.keys():
         fout = os.path.join(path, feature + '.zarr')
         if not os.path.exists(fout) or mode == 'w':
@@ -15,8 +55,8 @@ def xarray2zarr(xds, path, mode='a'):
         else:
             xds_existing = xr.open_zarr(fout, group='original')
             if xds_existing.datetime[0] > xds.datetime[0] or xds_existing.datetime[-1] > xds.datetime[-1]:
-                xds_new = xr.merge([xds_existing[feature], xds[feature]])
-                xds_new.to_zarr(fout, group='original', mode='w')
+                xda_new = merge_arrays(xds_existing[feature], xds[feature])
+                xda_new.to_zarr(fout, group='original', mode='w')
             else:
                 try:
                     overlap = xds_existing.datetime.where(
@@ -34,9 +74,5 @@ def xarray2zarr(xds, path, mode='a'):
                     msg += "Attempting to merge the two datasets."
                     logger.error(msg)
                     # remove duplicate datetime entries
-                    xda_existing = xds_existing[feature].drop_duplicates(
-                        'datetime', keep='last')
-                    xda_new = xds[feature].drop_duplicates(
-                        'datetime', keep='last')
-                    xda_new = xda_new.combine_first(xda_existing)
+                    xda_new = merge_arrays(xds_existing[feature], xds[feature])
                     xda_new.to_zarr(fout, group='original', mode='w')