fix: Include dtype information in the Partitioning configuration.

Author: Thomas Zilio

zcollection/partitioning/abc.py

@@ -106,7 +106,6 @@ def unique_and_check_monotony(arr: ArrayLike) -> tuple[NDArray, NDArray]:
 
     Args:
         arr: Array of elements.
-        is_delayed: If True, the array is delayed.
     Returns:
         Tuple of unique elements and their indices.
     """
@@ -331,12 +330,13 @@ def get_config(self) -> dict[str, Any]:
         Returns:
             The configuration of the partitioning scheme.
         """
-        config: dict[str, str | None] = {'id': self.ID}
+        config: dict[str, str | tuple[str, ...] | None] = {'id': self.ID}
         slots: Generator[tuple[str, ...]] = (getattr(
             _class, '__slots__',
             ()) for _class in reversed(self.__class__.__mro__))
         config.update((attr, getattr(self, attr)) for _class in slots
                       for attr in _class if not attr.startswith('_'))
+        config['dtype'] = self._dtype
         return config
 
     @classmethod

zcollection/partitioning/date.py

@@ -217,3 +217,10 @@ def decode(
         py_datetime: datetime.datetime = datetime64.astype('M8[s]').item()
         return tuple((UNITS[ix], getattr(py_datetime, self._attrs[ix]))
                      for ix in self._index)
+
+    def get_config(self) -> dict[str, Any]:
+        config = super().get_config()
+
+        # dtype are automatically computed by this partitioning
+        config.pop('dtype')
+        return config

zcollection/partitioning/tests/test_date.py

@@ -166,23 +166,43 @@ def test_construction() -> None:
         Date(('dates', ), 'W')
 
 
-def test_config():
+RESOLUTION_DTYPE_TEST_SET = [
+    ('Y', (('year', 'uint16'), )),
+    ('M', (('year', 'uint16'), ('month', 'uint8'))),
+    ('D', (('year', 'uint16'), ('month', 'uint8'), ('day', 'uint8'))),
+    ('h', (('year', 'uint16'), ('month', 'uint8'), ('day', 'uint8'),
+           ('hour', 'uint8'))),
+    ('m', (('year', 'uint16'), ('month', 'uint8'), ('day', 'uint8'),
+           ('hour', 'uint8'), ('minute', 'uint8'))),
+    ('s', (('year', 'uint16'), ('month', 'uint8'), ('day', 'uint8'),
+           ('hour', 'uint8'), ('minute', 'uint8'), ('second', 'uint8')))
+]
+
+
+@pytest.mark.parametrize('resolution, dtype', RESOLUTION_DTYPE_TEST_SET)
+def test_config(resolution, dtype):
     """Test the configuration of the Date class."""
-    partitioning = Date(('dates', ), 'D')
-    assert partitioning.dtype() == (('year', 'uint16'), ('month', 'uint8'),
-                                    ('day', 'uint8'))
+    partitioning = Date(variables=('dates', ), resolution=resolution)
+    assert partitioning.dtype() == dtype
+
     config = partitioning.get_config()
-    partitioning = get_codecs(config)
-    assert isinstance(partitioning, Date)
+    other = get_codecs(config)
+
+    assert isinstance(other, Date)
+    assert other.variables == ('dates', )
+    assert other.dtype() == dtype
 
 
-def test_pickle():
+@pytest.mark.parametrize('resolution, dtype', RESOLUTION_DTYPE_TEST_SET)
+def test_pickle(resolution, dtype):
     """Test the pickling of the Date class."""
-    partitioning = Date(('dates', ), 'D')
+    partitioning = Date(('dates', ), resolution=resolution)
     other = pickle.loads(pickle.dumps(partitioning))
+
     assert isinstance(other, Date)
-    assert other.resolution == 'D'
+    assert other.resolution == resolution
     assert other.variables == ('dates', )
+    assert other.dtype() == dtype
 
 
 @pytest.mark.parametrize('delayed', [False, True])

zcollection/partitioning/tests/test_sequence.py

@@ -113,20 +113,32 @@ def test_split_dataset(
         list(partitioning.split_dataset(zds, 'num_lines'))
 
 
-def test_config() -> None:
+VARIABLES_DTYPE_TEST_SET = [(('a', ), None), (('a', ), ('uint8', )),
+                            (('a', 'b'), None),
+                            (('a', 'b'), ('int8', 'int16'))]
+
+
+@pytest.mark.parametrize('variables, dtype', VARIABLES_DTYPE_TEST_SET)
+def test_config(variables, dtype) -> None:
     """Test the configuration of the Sequence class."""
-    partitioning = Sequence(('cycle_number', 'pass_number'))
+    partitioning = Sequence(variables=variables, dtype=dtype)
+
     config = partitioning.get_config()
-    partitioning = get_codecs(config)  # type: ignore[assignment]
-    assert isinstance(partitioning, Sequence)
+    other = get_codecs(config)  # type: ignore[assignment]
 
+    assert isinstance(other, Sequence)
+    assert other.dtype() == partitioning.dtype()
 
-def test_pickle() -> None:
+
+@pytest.mark.parametrize('variables, dtype', VARIABLES_DTYPE_TEST_SET)
+def test_pickle(variables, dtype) -> None:
     """Test the pickling of the Date class."""
-    partitioning = Sequence(('cycle_number', 'pass_number'))
+    partitioning = Sequence(variables=variables, dtype=dtype)
+
     other = pickle.loads(pickle.dumps(partitioning))
+
     assert isinstance(other, Sequence)
-    assert other.variables == ('cycle_number', 'pass_number')
+    assert other.dtype() == partitioning.dtype()
 
 
 # pylint: disable=protected-access