Experiments on a multi-file reader

satpy/readers/yaml_reader.py

@@ -1,4 +1,4 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 # Copyright (c) 2016-2022 Satpy developers
 #
@@ -23,7 +23,7 @@
 import os
 import warnings
 from abc import ABCMeta, abstractmethod
-from collections import OrderedDict, deque
+from collections import deque
 from contextlib import suppress
 from fnmatch import fnmatch
 from weakref import WeakValueDictionary
@@ -121,6 +121,11 @@
     return config
 
 
+def remove_duplicates(elements):
+    """Remove duplicates from a list while retaining order."""
+    return list(dict.fromkeys(elements))
+
+
 class AbstractYAMLReader(metaclass=ABCMeta):
     """Base class for all readers that use YAML configuration files.
 
@@ -278,7 +283,7 @@
         See `satpy.readers.get_key` for more information about kwargs.
 
         """
-        return get_key(key, self.all_ids.keys(), **kwargs)
+        return get_key(key, self.all_dataset_ids, **kwargs)
 
     def load_ds_ids_from_config(self):
         """Get the dataset ids from the config."""
@@ -359,7 +364,7 @@
                  filter_filenames=True,
                  **kwargs):
         """Set up initial internal storage for loading file data."""
-        super(FileYAMLReader, self).__init__(config_dict)
+        super().__init__(config_dict)
 
         self.file_handlers = {}
         self.available_ids = {}
@@ -605,7 +610,7 @@
 
     def create_filehandlers(self, filenames, fh_kwargs=None):
         """Organize the filenames into file types and create file handlers."""
-        filenames = list(OrderedDict.fromkeys(filenames))
+        filenames = remove_duplicates(filenames)
         logger.debug("Assigning to %s: %s", self.info["name"], filenames)
 
         self.info.setdefault("filenames", []).extend(filenames)
@@ -1164,7 +1169,7 @@
         produce the correct order.
 
         """
-        created_fhs = super(GEOSegmentYAMLReader, self).create_filehandlers(
+        created_fhs = super().create_filehandlers(
             filenames, fh_kwargs=fh_kwargs)
 
         # add "expected_segments" information
@@ -1542,3 +1547,50 @@
         ar = np.repeat(mod, n)
         ar[-remainder:] = mod + 1
         return ar.astype("int")
+
+
+class MultiFileYAMLReader(AbstractYAMLReader):
+    """Class for handling multiple files at once (same reader, same overpass/repeat cycle)."""
+
+    def start_time(self, *args, **kwargs):
+        """Get the start time of the data."""
+        pass
+
+    def end_time(self, *args, **kwargs):
+        """Get the end time of the data."""
+        pass
+
+    def filter_selected_filenames(self, *args, **kwargs):
+        """Filter filenames."""
+        pass
+
+    def create_filehandlers(self, files, **kwargs):
+        """Create file handlers."""
+        self.assign_storage_items(files)
+
+    def assign_storage_items(self, files):
+        """Assign storage items."""
+        self.json_file, self.data_file = files
+        if str(self.json_file).endswith("npz"):
+            self.json_file, self.data_file = self.data_file, self.json_file
+
+    @property
+    def available_dataset_ids(self):
+        """Generate the available dataset ids."""
+        return [DataID(default_id_keys_config, name="chanel_5", resolution=400)]
+
+    @property
+    def all_dataset_ids(self):
+        """Generate all the dataset ids."""
+        return [DataID(default_id_keys_config, name="chanel_5", resolution=400)]
+
+    def load(self, dataset_keys, **kwargs):
+        """Load the data."""
+        import json
+        with open(self.json_file, "r") as fp:
+            metadata = json.load(fp)
+        data_arrays = DatasetDict()
+        for dsid in dataset_keys:
+            data = np.load(self.data_file)[dsid["name"]]
+            data_arrays[dsid] = xr.DataArray(data, attrs=metadata)
+        return data_arrays

satpy/tests/test_multi_file_yaml_reader.py

@@ -0,0 +1,56 @@
+"""Tests for a multi file reader."""
+
+import numpy as np
+import pytest
+
+from satpy.readers.yaml_reader import MultiFileYAMLReader
+
+yaml_reader_config = {
+  "reader": {
+    "name": "json_npz",
+    "short_name": "Simple JSON and npz reader",
+    "sensors": ["vvhrr"],
+    "reader": MultiFileYAMLReader
+  },
+  "file_types": {"metadata": {"file_patterns": ["{start_time:%Y%m%dT%H%M%S}_vvhrr.json"]},
+                 "data": {"file_patterns": ["{start_time:%Y%m%dT%H%M%S}_vvhrr.data.npz"]}}
+                 }
+
+
+@pytest.fixture()
+def vvhrr_files(tmp_path):
+    """Create a fake file."""
+    import json
+    metadata = dict(sensor="vvhrr")
+    data = np.zeros((10, 10))
+    json_file = tmp_path / "20231211T111111_vvhrr.json"
+    data_file = tmp_path / "20231211T111111_vvhrr.data.npz"
+    with open(json_file, "w") as fp:
+        json.dump(metadata, fp)
+    np.savez(data_file, chanel_5=data)
+    return json_file, data_file
+
+
+def test_read_using_yaml_reader_interface(vvhrr_files):
+    """Test loading from storage using the YAMLReader interface."""
+    reader = MultiFileYAMLReader(yaml_reader_config)
+    reader.assign_storage_items(vvhrr_files)
+    from satpy.dataset.dataid import DataID, default_id_keys_config
+    dataarray_key = DataID(default_id_keys_config, name="chanel_5", resolution=400)
+    res = reader.load([dataarray_key])[dataarray_key]
+    expected = np.zeros((10, 10))
+    np.testing.assert_allclose(res, expected)
+    assert res.attrs["sensor"] == "vvhrr"
+
+def test_read_using_scene_interface(vvhrr_files, tmp_path):
+    """Test the reader interface."""
+    import yaml
+
+    from satpy import Scene, config
+    config_dir = tmp_path / "readers"
+    config_dir.mkdir()
+    with open(config_dir / "json_npz.yaml", "w") as fd:
+        fd.write(yaml.dump(yaml_reader_config))
+    with config.set(config_path=[tmp_path]):
+        scn = Scene(vvhrr_files, reader="json_npz")
+        scn.load(["chanel_5"])