Refactor offline I/O

This is a rebased branch of !24 (closed) and ready to merge

km3io/tools.py

@@ -19,6 +19,21 @@ class cached_property:
         return prop
 
 
+def _unfold_indices(obj, indices):
+    """Unfolds an index chain and returns the corresponding item"""
+    original_obj = obj
+    for depth, idx in enumerate(indices):
+        try:
+            obj = obj[idx]
+        except IndexError:
+            print(
+                "IndexError while accessing an item from '{}' at depth {} ({}) "
+                "using the index chain {}".format(repr(original_obj), depth,
+                                                  idx, indices))
+            raise
+    return obj
+
+
 BranchMapper = namedtuple(
     "BranchMapper",
     ['name', 'key', 'extra', 'exclude', 'update', 'attrparser', 'flat'])
@@ -29,17 +44,19 @@ class Branch:
     def __init__(self,
                  tree,
                  mapper,
-                 index=None,
+                 index_chain=None,
                  subbranchmaps=None,
                  keymap=None):
         self._tree = tree
         self._mapper = mapper
-        self._index = index
+        self._index_chain = [] if index_chain is None else index_chain
         self._keymap = None
         self._branch = tree[mapper.key]
         self._subbranches = []
         self._subbranchmaps = subbranchmaps
 
+        self._iterator_index = 0
+
         if keymap is None:
             self._initialise_keys()  #
         else:
@@ -49,7 +66,7 @@ class Branch:
             for mapper in subbranchmaps:
                 subbranch = self.__class__(self._tree,
                                            mapper=mapper,
-                                           index=self._index)
+                                           index_chain=self._index_chain)
                 self._subbranches.append(subbranch)
         for subbranch in self._subbranches:
             setattr(self, subbranch._mapper.name, subbranch)
@@ -57,8 +74,8 @@ class Branch:
     def _initialise_keys(self):
         """Create the keymap and instance attributes for branch keys"""
         # TODO: this could be a cached property
-        keys = set(k.decode('utf-8') for k in self._branch.keys()) - set(
-            self._mapper.exclude)
+        keys = set(k.decode('utf-8')
+                   for k in self._branch.keys()) - set(self._mapper.exclude)
         self._keymap = {
             **{self._mapper.attrparser(k): k
                for k in keys},
@@ -86,42 +103,46 @@ class Branch:
     def __getkey__(self, key):
         out = self._branch[self._keymap[key]].lazyarray(
             basketcache=BASKET_CACHE)
-        if self._index is not None:
-            out = out[self._index]
-        return out
+        return _unfold_indices(out, self._index_chain)
 
     def __getitem__(self, item):
         """Slicing magic"""
-        if isinstance(item, (int, slice)):
-            return self.__class__(self._tree,
-                                  self._mapper,
-                                  index=item,
-                                  keymap=self._keymap,
-                                  subbranchmaps=self._subbranchmaps)
-
-        if isinstance(item, tuple):
-            return self[item[0]][item[1]]
-
         if isinstance(item, str):
             return self.__getkey__(item)
 
         return self.__class__(self._tree,
                               self._mapper,
-                              index=np.array(item),
+                              index_chain=self._index_chain + [item],
                               keymap=self._keymap,
                               subbranchmaps=self._subbranchmaps)
 
     def __len__(self):
-        if self._index is None:
+        if not self._index_chain:
             return len(self._branch)
-        elif isinstance(self._index, int):
+        elif isinstance(self._index_chain[-1], int):
             return 1
         else:
-            return len(self._branch[self._keymap['id']].lazyarray(
-                basketcache=BASKET_CACHE)[self._index])
+            return len(
+                _unfold_indices(
+                    self._branch[self._keymap['id']].lazyarray(
+                        basketcache=BASKET_CACHE), self._index_chain))
+
+    def __iter__(self):
+        self._iterator_index = 0
+        return self
+
+    def __next__(self):
+        idx = self._iterator_index
+        self._iterator_index += 1
+        if idx >= len(self):
+            raise StopIteration
+        return self[idx]
 
     def __str__(self):
-        return "Number of elements: {}".format(len(self._branch))
+        length = len(self)
+        return "{} ({}) with {} element{}".format(self.__class__.__name__,
+                                               self._mapper.name, length,
+                                               's' if length > 1 else '')
 
     def __repr__(self):
         length = len(self)