fix: save_metadata to support both multiple datasets with many arrays and old set up with no arrays (#239)

## Description
- PR to fix #238

Add test to write metadata with no arrays, and when arrays dict contains
multiple entries (simulating multiple-dataset set up)

## What problem does this change solve?
Previous changes broke support for old set ups where supporting arrays
were not defined.

## What issue or task does this change relate to?
[<!-- link to Issue Number
-->](#238)

##  Additional notes ##
<!-- Include any additional information, caveats, or considerations that
the reviewer should be aware of. -->

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

Author: anaprietonem

src/anemoi/utils/checkpoints.py

@@ -141,10 +141,13 @@ def load_supporting_arrays(zipf: zipfile.ZipFile, entries: dict) -> dict:
 
     supporting_arrays = {}
     for key, entry in entries.items():
-        supporting_arrays[key] = np.frombuffer(
-            zipf.read(entry["path"]),
-            dtype=entry["dtype"],
-        ).reshape(entry["shape"])
+        if isinstance(entry, dict) and not set(entry.keys()) == set(["path", "shape", "dtype"]):
+            supporting_arrays[key] = load_supporting_arrays(zipf, entry)
+        else:
+            supporting_arrays[key] = np.frombuffer(
+                zipf.read(entry["path"]),
+                dtype=entry["dtype"],
+            ).reshape(entry["shape"])
     return supporting_arrays
 
 
@@ -168,11 +171,14 @@ def _get_supporting_arrays_paths(directory: str, folder: str, supporting_arrays:
 
 def _write_array_to_bytes(array: dict | np.ndarray, name: str, entry: dict, zipf: zipfile.ZipFile) -> None:
     """Write a supporting array to bytes in a zip file."""
+    if array is None:
+        return
+
     if isinstance(array, dict):
         for sub_name, sub_array in array.items():
-            _write_array_to_bytes(sub_array, sub_name, entry[sub_name], zipf)
-        return None
-
+            sub_entry = entry.get(sub_name, {})
+            _write_array_to_bytes(sub_array, sub_name, sub_entry, zipf)
+        return
     LOG.info(
         "Saving supporting array `%s` to %s (shape=%s, dtype=%s)",
         name,

tests/test_checkpoints.py

@@ -187,21 +187,32 @@ def test_remove_metadata_integration(self, sample_checkpoint):
 
         assert not has_metadata(sample_checkpoint, name="metadata.json")
 
-    def test_metadata_with_arrays_roundtrip(self, sample_checkpoint):
+    @pytest.mark.parametrize(
+        "arrays",
+        [
+            {"test_array": np.array([1, 2, 3, 4, 5])},
+            {
+                "test_array": {"test_array": np.array([10, 20, 30])},
+            },
+        ],
+    )
+    def test_metadata_with_arrays_roundtrip(self, sample_checkpoint, arrays):
         """Test complete roundtrip with supporting arrays."""
         # First remove existing metadata
         remove_metadata(sample_checkpoint, name="metadata.json")
 
         # Add metadata with arrays
         metadata = {"version": "1.0", "test": True}
-        arrays = {"test_array": np.array([1, 2, 3, 4, 5])}
 
         save_metadata(sample_checkpoint, metadata, supporting_arrays=arrays, name="metadata.json")
 
         # Load and verify
         loaded_metadata, loaded_arrays = load_metadata(sample_checkpoint, supporting_arrays=True, name="metadata.json")
         assert loaded_metadata["test"] is True
-        np.testing.assert_array_equal(loaded_arrays["test_array"], arrays["test_array"])
+        if isinstance(loaded_arrays["test_array"], dict):
+            np.testing.assert_array_equal(loaded_arrays["test_array"]["test_array"], arrays["test_array"]["test_array"])
+        else:
+            np.testing.assert_array_equal(loaded_arrays["test_array"], arrays["test_array"])
 
         # Edit with _edit_metadata
         def update_callback(file_path):
@@ -217,4 +228,31 @@ def update_callback(file_path):
         final_metadata, final_arrays = load_metadata(sample_checkpoint, supporting_arrays=True, name="metadata.json")
         assert final_metadata["edited"] is True
         assert final_metadata["test"] is True
-        np.testing.assert_array_equal(final_arrays["test_array"], arrays["test_array"])
+        if isinstance(loaded_arrays["test_array"], dict):
+            np.testing.assert_array_equal(final_arrays["test_array"]["test_array"], arrays["test_array"]["test_array"])
+        else:
+            np.testing.assert_array_equal(final_arrays["test_array"], arrays["test_array"])
+
+    def test_metadata_no_arrays(self, sample_checkpoint):
+        """Test without supporting arrays."""
+        # First remove existing metadata
+        remove_metadata(sample_checkpoint, name="metadata.json")
+
+        # Add metadata with arrays
+        metadata = {"version": "1.0", "test": True}
+
+        save_metadata(sample_checkpoint, metadata, supporting_arrays=None, name="metadata.json")
+
+        # Load and verify
+        loaded_metadata = load_metadata(sample_checkpoint, supporting_arrays=None, name="metadata.json")
+        assert loaded_metadata["test"] is True
+
+        # Edit with _edit_metadata
+        def update_callback(file_path):
+            with open(file_path, "r") as f:
+                data = json.load(f)
+            data["edited"] = True
+            with open(file_path, "w") as f:
+                json.dump(data, f)
+
+        _edit_metadata(sample_checkpoint, "metadata.json", update_callback)