[DEM-NNN] Load data from subtag (#100)

* add method to return data from subtag

Co-authored-by: Pieter Eendebak <[email protected]>

src/qilib/utils/storage/interface.py

@@ -21,11 +21,41 @@
 import logging
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Callable, Optional, Union
+from collections.abc import Sequence as SequenceBaseClass
+from typing import Any, Callable, Optional, Union, Sequence
 
 from qilib.utils.type_aliases import TagType
 
 
+class _LazySequence(SequenceBaseClass):  # type: ignore
+    def __init__(self, length: int, item_getter: Callable[[int], Any]):
+        """ Convert a length and method to retrieve an indexed item into a sequence with lazy evaluation
+
+        Args:
+            length: Length of the sequence to be represented
+            item_getter: Method to retrieve an item at the specified index
+        """
+        super(SequenceBaseClass, self).__init__()
+        self._length = length
+        self._item_getter = item_getter
+
+    def __repr__(self) -> str:
+        classname = ".".join([self.__module__, self.__class__.__qualname__])
+        return f'<{classname} at %x{id(self)} length {self._length} >'
+
+    def __len__(self) -> int:
+        r = self._length
+        return r
+
+    def __getitem__(self, index: Union[int, slice]) -> Any:
+        if isinstance(index, slice):
+            slice_range = range(*index.indices(len(self)))
+            return (self._item_getter(i) for i in slice_range)
+        else:
+            r = self._item_getter(index)
+            return r
+
+
 class NoDataAtKeyError(Exception):
     """ Raised when trying to get data from a node or leave can not be found."""
 
@@ -193,11 +223,12 @@ def get_latest_subtag(self, tag: TagType) -> Optional[TagType]:
         pass
 
     @abstractmethod
-    def list_data_subtags(self, tag: TagType) -> TagType:
+    def list_data_subtags(self, tag: TagType, limit: int = 0) -> TagType:
         """ List available result tags of at tag.
 
         Args:
             tag: hdf5 tag
+            limit: Maximum number of results to generate. If 0 then return all results
         Returns:
             results: List of child tags for tag
 
@@ -212,6 +243,22 @@ def list_data_subtags(self, tag: TagType) -> TagType:
         """
         pass
 
+    def load_data_from_subtag(self, tag: TagType, limit: int = 0) -> Sequence[Any]:
+        """ Return all results under the specified tag
+
+        Args:
+            tag: Tag to search for results
+            limit: Maximum number of results to generate. If 0 then return all results
+        Returns:
+            Generator for all the results
+        """
+        subtags = self.list_data_subtags(tag, limit=limit)
+
+        def item_getter(index: int) -> Any:
+            subtag = subtags[index]
+            return self.load_data(tag+[subtag])
+        return _LazySequence(len(subtags), item_getter)
+
     @abstractmethod
     def search(self, query: str) -> Any:
         """ Future implementation of query interface """

src/tests/unittests/utils/storage/test_storage_interface.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import patch
 
-from qilib.utils.storage.interface import StorageInterface
+from qilib.utils.storage.interface import StorageInterface, _LazySequence
 
 
 class TestStorage(unittest.TestCase):
@@ -32,3 +32,10 @@ def test_abc_dummy_tests(self):
             storage_interface.load_individual_data(None, None)
             storage_interface.update_individual_data(None, None, None)
             self.assertRaises(NotImplementedError, storage_interface.search, None)
+
+    def test_LazySequence(self):
+        getter = lambda i: i
+        lazy_list = _LazySequence(10, getter)
+        self.assertEqual(lazy_list[2], 2)
+        self.assertEqual(list(lazy_list[6:]), [6,7,8,9])
+        self.assertIn('length 10', repr(lazy_list) )

src/tests/unittests/utils/storage/test_storage_memory.py

@@ -158,3 +158,10 @@ def test_update_individual_data(self):
 
         self.assertRaises(NodeDoesNotExistsError, self.storage.update_individual_data, 42, ['data', str(3000)], 'a')
         self.assertRaises(NodeDoesNotExistsError, self.storage.update_individual_data, 42, ['data3000'], 'a')
+
+    def test_load_data_from_subtag(self):
+        storage_interface = StorageMemory('test')
+        for ii in range(4):
+            storage_interface.save_data(ii, ['s', f's{ii}'])        
+        self.assertEqual(list(storage_interface.load_data_from_subtag(['s'])), [0,1,2,3])
+