Merge pull request #54 from BAMresearch/improve_io_sources

Improve io sources

Author: toqduj

src/modacor/io/hdf/hdf_loader.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any
 
 __coding__ = "utf-8"
 __author__ = "Tim Snow, Brian R. Pauw"
@@ -35,7 +35,7 @@ class HDFLoader(IoSource):
     _static_metadata_cache: dict[str, Any] = None
 
     def __init__(self, source_reference: str, logging_level=WARNING, resource_location: Path | str | None = None):
-        super().__init__(source_reference)
+        super().__init__(source_reference=source_reference)
         self.logger = MessageHandler(level=logging_level, name="HDFLoader")
         self._file_path = Path(resource_location) if resource_location is not None else None
         # self._file_reference = None  # let's not leave open file references lying around if we can help it.
@@ -44,6 +44,7 @@ def __init__(self, source_reference: str, logging_level=WARNING, resource_locati
         self._file_datasets_dtypes = {}
         self._data_cache = {}
         self._static_metadata_cache = {}
+        self._preload()  # load the HDF5 file structure immediately so we have some information, but not the data
 
     def _preload(self):
         assert self._file_path.is_file(), self.logger.error(f"HDF5 file {self._file_path} does not exist.")

src/modacor/io/io_registry.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+import attrs
+
 __coding__ = "utf-8"
 __authors__ = ["Malte Storm", "Brian R. Pauw"]  # add names to the list as appropriate
 __copyright__ = "Copyright 2025, The MoDaCor team"
@@ -60,9 +62,9 @@ class IoSource:
             slices can be separated by double semicolon ';;'.
     """
 
-    type_reference = "IoSource"
-
     configuration: dict[str, Any] = field(factory=default_config)
+    source_reference: str = field(default="", converter=str, validator=attrs.validators.instance_of(str))
+    type_reference: str = "IoSource"
 
     def get_data(self, data_key: str, load_slice: Optional[ArraySlice] = None) -> np.ndarray:
         """

src/modacor/io/io_source.py

@@ -1,39 +1,25 @@
 # SPDX-License-Identifier: BSD-3-Clause
-# Copyright 2025 MoDaCor Authors
-#
-# Redistribution and use in source and binary forms, with or without modification,
-# are permitted provided that the following conditions are met:
-# 1. Redistributions of source code must retain the above copyright notice, this
-#    list of conditions and the following disclaimer.
-# 2. Redistributions in binary form must reproduce the above copyright notice,
-#    this list of conditions and the following disclaimer in the documentation
-#    and/or other materials provided with the distribution.
-# 3. Neither the name of the copyright holder nor the names of its contributors
-#    may be used to endorse or promote products derived from this software without
-#    specific prior written permission.
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND
-# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
-# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
-# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
-# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-__license__ = "BSD-3-Clause"
-__copyright__ = "Copyright 2025 MoDaCor Authors"
-__status__ = "Alpha"
+# /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+__coding__ = "utf-8"
+__author__ = "Brian R. Pauw"
+__copyright__ = "Copyright 2025, The MoDaCor team"
+__date__ = "06/06/2025"
+__status__ = "Development"  # "Development", "Production"
+# end of header and standard imports
+
 __all__ = ["IoSources"]
 
 
-from typing import Any
+from typing import Any, Optional
 
 import numpy as np
 from attrs import define, field
 
-from modacor.io.io_source import IoSource
+from modacor.io.io_source import ArraySlice, IoSource
 
 
 @define
@@ -47,21 +33,24 @@ class IoSources:
 
     defined_sources: dict[str, IoSource] = field(factory=dict)
 
-    def register_source(self, source_reference: str, source: IoSource):
+    def register_source(self, source: IoSource, source_reference: str | None = None):
         """
-        Register a new source class with the given name.
+        Register a new source class with the given name. If no source_reference is provided, the
+        source's own source_reference attribute will be used.
 
         Parameters
         ----------
-        source_reference : str
-            The reference name of the source to register.
         source : IoSource
             The class of the source to register.
+        source_reference : str
+            The reference name of the source to register.
         """
-        if not isinstance(source_reference, str):
-            raise TypeError("source_name must be a string")
         if not isinstance(source, IoSource):
             raise TypeError("source_class must be a subclass of IoSource")
+        if source_reference is None:
+            source_reference = source.source_reference
+        if not isinstance(source_reference, str):
+            raise TypeError("source_name must be a string")
         if source_reference in self.defined_sources:
             raise ValueError(f"Source {source_reference} already registered.")
         self.defined_sources[source_reference] = source
@@ -109,7 +98,7 @@ def split_data_reference(self, data_reference: str) -> tuple[str, str]:
             )
         return _split[0], _split[1]
 
-    def get_data(self, data_reference: str, index: int) -> np.ndarray:
+    def get_data(self, data_reference: str, load_slice: Optional[ArraySlice] = ...) -> np.ndarray:
         """
         Get data from the specified source using the provided data key.
 
@@ -120,8 +109,8 @@ def get_data(self, data_reference: str, index: int) -> np.ndarray:
         ----------
         data_reference : str
             The reference name of the source to access.
-        index : int
-            The index to access the data.
+        load_slice : Optional[ArraySlice]
+            A slice or tuple of slices to apply to the data. If None or ellipsis, the entire data is returned.
 
         Returns
         -------
@@ -130,9 +119,9 @@ def get_data(self, data_reference: str, index: int) -> np.ndarray:
         """
         _source_ref, _data_key = self.split_data_reference(data_reference)
         _source = self.get_source(_source_ref)
-        return _source.get_data(index, _data_key)
+        return _source.get_data(_data_key, load_slice=load_slice)
 
-    def get_data_shape(self, data_reference: str, index: int) -> np.ndarray:
+    def get_data_shape(self, data_reference: str) -> np.ndarray:
         """
         Get data from the specified source using the provided data key.
 
@@ -143,8 +132,6 @@ def get_data_shape(self, data_reference: str, index: int) -> np.ndarray:
         ----------
         data_reference : str
             The reference name of the source to access.
-        index : int
-            The index to access the data.
 
         Returns
         -------
@@ -153,9 +140,9 @@ def get_data_shape(self, data_reference: str, index: int) -> np.ndarray:
         """
         _source_ref, _data_key = self.split_data_reference(data_reference)
         _source = self.get_source(_source_ref)
-        return _source.get_data_shape(index, _data_key)
+        return _source.get_data_shape(_data_key)
 
-    def get_data_dtype(self, data_reference: str, index: int) -> np.ndarray:
+    def get_data_dtype(self, data_reference: str) -> np.ndarray:
         """
         Get data from the specified source using the provided data key.
 
@@ -166,8 +153,6 @@ def get_data_dtype(self, data_reference: str, index: int) -> np.ndarray:
         ----------
         data_reference : str
             The reference name of the source to access.
-        index : int
-            The index to access the data.
 
         Returns
         -------
@@ -176,9 +161,9 @@ def get_data_dtype(self, data_reference: str, index: int) -> np.ndarray:
         """
         _source_ref, _data_key = self.split_data_reference(data_reference)
         _source = self.get_source(_source_ref)
-        return _source.get_data_dtype(index, _data_key)
+        return _source.get_data_dtype(_data_key)
 
-    def get_data_attributes(self, data_reference: str, index: int) -> np.ndarray:
+    def get_data_attributes(self, data_reference: str) -> np.ndarray:
         """
         Get data from the specified source using the provided data key.
 
@@ -199,7 +184,7 @@ def get_data_attributes(self, data_reference: str, index: int) -> np.ndarray:
         """
         _source_ref, _data_key = self.split_data_reference(data_reference)
         _source = self.get_source(_source_ref)
-        return _source.get_data_attributes(index, _data_key)
+        return _source.get_data_attributes(_data_key)
 
     def get_static_metadata(self, data_reference: str) -> Any:
         """

src/modacor/io/io_sources.py

@@ -54,13 +54,14 @@ class YAMLLoader(IoSource):
     _static_metadata_cache: dict[str, Any] = None
 
     def __init__(self, source_reference: str, logging_level=WARNING, resource_location: Path | str | None = None):
-        super().__init__(source_reference)
+        super().__init__(source_reference=source_reference)
         self.logger = MessageHandler(level=logging_level, name="YAMLLoader")
         self._file_path = Path(resource_location) if resource_location is not None else None
         self._file_datasets = []
         self._file_datasets_shapes = {}
         self._data_cache = {}  # for values that are float
         self._static_metadata_cache = {}  # for other elements such as strings and tags
+        self._preload()  # load the yaml data immediately
 
     def _preload(self) -> None:
         """

src/modacor/io/yaml/yaml_loader.py

@@ -26,18 +26,20 @@ class TestHDFLoader(unittest.TestCase):
     """Testing class for modacor/io/hdf/hdf_loader.py"""
 
     def setUp(self):
-        self.test_hdf_loader = HDFLoader(source_reference="Test Data")
         self.temp_file_handle = tempfile.NamedTemporaryFile(delete=False, delete_on_close=False)
         self.temp_file_path = self.temp_file_handle.name
         self.temp_file_handle.close()
         self.temp_dataset_name = "dataset"
         self.temp_dataset_shape = (10, 2)
-        self.temp_hdf_file = h5py.File(self.temp_file_path, "w")
-        self.temp_hdf_file[self.temp_dataset_name] = np.zeros(self.temp_dataset_shape)
-        self.temp_file_handle.close()
+        with h5py.File(self.temp_file_path, "w") as hdf_file:
+            hdf_file.create_dataset(
+                self.temp_dataset_name, data=np.zeros(self.temp_dataset_shape), dtype="float64", compression="gzip"
+            )
+
+        self.test_hdf_loader = HDFLoader(source_reference="Test Data", resource_location=self.temp_file_path)
 
     def tearDown(self):
-        self.test_h5_loader = None
+        self.test_hdf_loader = None
         self.test_file_path = None
         self.test_dataset_name = None
         self.test_dataset_shape = None