Merge pull request #26 from bcdev/forman-25-sizes_instead_of_dims

Use Dataset.sizes instead of .dims

Author: forman

CHANGES.md

@@ -1,5 +1,7 @@
 ## Version 0.2.1 (in development)
 
+* Using `sizes` instead of `dims` attribute of `xarray.Dataset` in implementation 
+  code. [#25] 
 * Enhanced documentation including docstrings of several Python API objects.
 
 

tests/fsutil/test_transaction.py

@@ -92,10 +92,10 @@ def create_test_folder(rollback_cb: Callable):
                     )
 
                 if fail:
-                    raise OSError("disk full")
+                    raise OSError("disk full (this is a test!)")
         except OSError as e:
             if fail:
-                self.assertEqual("disk full", f"{e}")
+                self.assertEqual("disk full (this is a test!)", f"{e}")
             else:
                 raise
 

tests/test_api.py

@@ -25,6 +25,6 @@ def test_some_slices(self):
         slices = [make_test_dataset(), make_test_dataset(), make_test_dataset()]
         zappend(slices, target_dir=target_dir)
         ds = xr.open_zarr(target_dir)
-        self.assertEqual({"time": 9, "y": 50, "x": 100}, ds.dims)
+        self.assertEqual({"time": 9, "y": 50, "x": 100}, ds.sizes)
         self.assertEqual({"chl", "tsm"}, set(ds.data_vars))
         self.assertEqual({"time", "y", "x"}, set(ds.coords))

tests/test_metadata.py

@@ -18,7 +18,7 @@ def test_dims_without_fixed_dims_given(self):
         ds = xr.Dataset({"a": xr.DataArray(np.zeros((2, 3, 4)), dims=("z", "y", "x"))})
         self.assertEqual(
             {"z": 2, "y": 3, "x": 4},
-            DatasetMetadata.from_dataset(ds, {"append_dim": "z"}).dims,
+            DatasetMetadata.from_dataset(ds, {"append_dim": "z"}).sizes,
         )
 
     def test_dims_with_fixed_dims_given(self):
@@ -27,7 +27,7 @@ def test_dims_with_fixed_dims_given(self):
             {"z": 2, "y": 3, "x": 4},
             DatasetMetadata.from_dataset(
                 ds, {"append_dim": "z", "fixed_dims": {"y": 3, "x": 4}}
-            ).dims,
+            ).sizes,
         )
 
     # noinspection PyMethodMayBeStatic
@@ -84,7 +84,7 @@ def test_add_missing_variables(self):
         self.assertEqual(
             {
                 "attrs": {},
-                "dims": {"time": 2, "x": 4, "y": 3},
+                "sizes": {"time": 2, "x": 4, "y": 3},
                 "variables": {
                     "a": {
                         "attrs": {},
@@ -133,7 +133,7 @@ def test_merge_variable_metadata(self):
         self.assertEqual(
             {
                 "attrs": {},
-                "dims": {"time": 2, "x": 4, "y": 3},
+                "sizes": {"time": 2, "x": 4, "y": 3},
                 "variables": {
                     "a": {
                         "attrs": {"title": "A", "units": "m/s"},
@@ -159,7 +159,7 @@ def test_move_variable_encoding_from_attrs(self):
         self.assertEqual(
             {
                 "attrs": {},
-                "dims": {"time": 2, "x": 4, "y": 3},
+                "sizes": {"time": 2, "x": 4, "y": 3},
                 "variables": {
                     "a": {
                         "attrs": {},
@@ -199,7 +199,7 @@ def test_variable_defaults(self):
         self.assertEqual(
             {
                 "attrs": {},
-                "dims": {"time": 2, "x": 4, "y": 3},
+                "sizes": {"time": 2, "x": 4, "y": 3},
                 "variables": {
                     "a": {
                         "attrs": {},
@@ -246,7 +246,7 @@ def test_variable_encoding_from_netcdf(self):
         self.assertEqual(
             {
                 "attrs": {},
-                "dims": {"time": 2, "x": 4, "y": 3},
+                "sizes": {"time": 2, "x": 4, "y": 3},
                 "variables": {
                     "a": {
                         "attrs": {},

tests/test_processor.py

@@ -28,7 +28,7 @@ def test_process_one_slice(self):
         ds = xr.open_zarr(
             target_dir.uri, storage_options=target_dir.storage_options, decode_cf=True
         )
-        self.assertEqual({"time": 1, "y": 10, "x": 20}, ds.dims)
+        self.assertEqual({"time": 1, "y": 10, "x": 20}, ds.sizes)
         self.assertEqual({"x", "y", "time", "chl", "tsm"}, set(ds.variables))
 
         self.assertEqual((20,), ds.x.encoding.get("chunks"))
@@ -53,7 +53,7 @@ def test_process_two_slices(self):
         )
 
         self.assertEqual({"x", "y", "time", "chl", "tsm"}, set(ds.variables))
-        self.assertEqual({"time": 2, "y": 10, "x": 20}, ds.dims)
+        self.assertEqual({"time": 2, "y": 10, "x": 20}, ds.sizes)
 
         self.assertEqual((20,), ds.x.encoding.get("chunks"))
         self.assertEqual((10,), ds.y.encoding.get("chunks"))
@@ -91,7 +91,7 @@ def test_process_many_slices_with_single_append_dim_chunk(self):
         )
 
         self.assertEqual({"x", "y", "time", "chl", "tsm"}, set(ds.variables))
-        self.assertEqual({"time": many, "y": 10, "x": 20}, ds.dims)
+        self.assertEqual({"time": many, "y": 10, "x": 20}, ds.sizes)
 
         self.assertEqual((20,), ds.x.encoding.get("chunks"))
         self.assertEqual((10,), ds.y.encoding.get("chunks"))
@@ -129,7 +129,7 @@ def test_process_two_slices_with_chunk_overlap(self):
         )
 
         self.assertEqual({"x", "y", "time", "chl", "tsm"}, set(ds.variables))
-        self.assertEqual({"time": 4, "y": 10, "x": 20}, ds.dims)
+        self.assertEqual({"time": 4, "y": 10, "x": 20}, ds.sizes)
 
         self.assertEqual((20,), ds.x.encoding.get("chunks"))
         self.assertEqual((10,), ds.y.encoding.get("chunks"))

tests/test_rollbackstore.py

@@ -196,7 +196,7 @@ def test_to_zarr(self):
         slice_1 = make_test_dataset(shape=(1, 50, 100), chunks=(1, 50, 50))
         # drop variables w.o. "time" dim
         slice_1 = slice_1.drop_vars(
-            [k for k, v in slice_1.variables.items() if "time" not in v.dims]
+            [k for k, v in slice_1.variables.items() if "time" not in v.sizes]
         )
         slice_1.attrs = {}
         for k, v in slice_1.variables.items():
@@ -247,7 +247,7 @@ def test_to_zarr(self):
         slice_2 = make_test_dataset(shape=(1, 50, 100), chunks=(1, 50, 50))
         # drop variables w.o. "time" dim
         slice_2 = slice_2.drop_vars(
-            [k for k, v in slice_2.variables.items() if "time" not in v.dims]
+            [k for k, v in slice_2.variables.items() if "time" not in v.sizes]
         )
         for k, v in slice_2.variables.items():
             v.encoding = {}
@@ -291,10 +291,12 @@ def test_to_zarr(self):
         )
 
     def assert_dataset_ok(
-        self, expected_dims: dict[str, int], expected_chunks: dict[str, tuple[int, ...]]
+        self,
+        expected_sizes: dict[str, int],
+        expected_chunks: dict[str, tuple[int, ...]],
     ):
         ds = xr.open_zarr(self.target_dir.uri)
-        self.assertEqual(expected_dims, ds.dims)
+        self.assertEqual(expected_sizes, ds.sizes)
         self.assertEqual(
             expected_chunks,
             {k: ds[k].encoding.get("chunks") for k in ds.variables.keys()},

zappend/metadata.py

@@ -91,28 +91,28 @@ def to_dict(self):
 class DatasetMetadata:
     def __init__(
         self,
-        dims: dict[str, int],
+        sizes: dict[str, int],
         variables: dict[str, VariableMetadata],
         attrs: dict[str, Any],
     ):
-        self.dims = dims
+        self.sizes = sizes
         self.variables = variables
         self.attrs = attrs
 
     def to_dict(self):
         return dict(
-            dims=self.dims,
+            sizes=self.sizes,
             variables={k: v.to_dict() for k, v in self.variables.items()},
             attrs=self.attrs,
         )
 
     def assert_compatible_slice(
         self, slice_metadata: "DatasetMetadata", append_dim: str
     ):
-        for dim_name, dim_size in self.dims.items():
-            if dim_name not in slice_metadata.dims:
+        for dim_name, dim_size in self.sizes.items():
+            if dim_name not in slice_metadata.sizes:
                 raise ValueError(f"Missing dimension" f" {dim_name!r} in slice dataset")
-            slice_dim_size = slice_metadata.dims[dim_name]
+            slice_dim_size = slice_metadata.sizes[dim_name]
             if dim_name != append_dim and dim_size != slice_dim_size:
                 raise ValueError(
                     f"Wrong size for dimension {dim_name!r}"
@@ -137,7 +137,7 @@ def assert_compatible_slice(
     def from_dataset(cls, dataset: xr.Dataset, config: dict[str, Any] | None = None):
         config = config or {}
 
-        dims = _get_effective_dims(
+        sizes = _get_effective_sizes(
             dataset,
             config.get("fixed_dims"),
             config.get("append_dim") or DEFAULT_APPEND_DIM,
@@ -152,10 +152,10 @@ def from_dataset(cls, dataset: xr.Dataset, config: dict[str, Any] | None = None)
 
         attrs = merge_configs(dataset.attrs, config.get("attrs") or {})
 
-        return DatasetMetadata(dims=dims, variables=variables, attrs=attrs)
+        return DatasetMetadata(sizes=sizes, variables=variables, attrs=attrs)
 
 
-def _get_effective_dims(
+def _get_effective_sizes(
     dataset: xr.Dataset,
     config_fixed_dims: dict[str, int] | None,
     config_append_dim: str,
@@ -170,19 +170,19 @@ def _get_effective_dims(
                 raise ValueError(
                     f"Fixed dimension {dim_name!r}" f" not found in dataset"
                 )
-            ds_dim_size = dataset.dims[dim_name]
+            ds_dim_size = dataset.sizes[dim_name]
             if fixed_dim_size != ds_dim_size:
                 raise ValueError(
                     f"Wrong size for fixed dimension {dim_name!r}"
                     f" in dataset: expected {fixed_dim_size},"
                     f" found {ds_dim_size}"
                 )
-    if config_append_dim not in dataset.dims:
+    if config_append_dim not in dataset.sizes:
         raise ValueError(
             f"Append dimension" f" {config_append_dim!r} not found in dataset"
         )
 
-    return {str(k): v for k, v in dataset.dims.items()}
+    return {str(k): v for k, v in dataset.sizes.items()}
 
 
 def _get_effective_variables(
@@ -248,12 +248,12 @@ def _get_effective_variables(
             if config_var_dims is None:
                 raise ValueError(f"Missing dimensions" f" of variable {var_name!r}")
             for dim in config_var_dims:
-                if dim not in dataset.dims:
+                if dim not in dataset.sizes:
                     raise ValueError(
                         f"Dimension {dim!r} of variable"
                         f" {var_name!r} not found in dataset"
                     )
-            config_var_def["shape"] = tuple(dataset.dims[k] for k in config_var_dims)
+            config_var_def["shape"] = tuple(dataset.sizes[k] for k in config_var_dims)
             encoding: dict | None = config_var_def.get("encoding")
             if encoding is None or encoding.get("dtype") is None:
                 raise ValueError(