Resolve "uproot4 integration"

Closes #58 (closed)

km3io/offline.py

@@ -79,7 +79,15 @@ class OfflineReader:
         "mc_tracks": "mc_trks",
     }
 
-    def __init__(self, f, index_chain=None, step_size=2000, keys=None, aliases=None, event_ctor=None):
+    def __init__(
+        self,
+        f,
+        index_chain=None,
+        step_size=2000,
+        keys=None,
+        aliases=None,
+        event_ctor=None,
+    ):
         """OfflineReader class is an offline ROOT file wrapper
 
         Parameters
@@ -178,8 +186,10 @@ class OfflineReader:
 
     def __getitem__(self, key):
         # indexing
-        if isinstance(key, (slice, int, np.int32, np.int64)):
-            if not isinstance(key, slice):
+        # TODO: maybe just propagate everything to awkward and let it deal
+        # with the type?
+        if isinstance(key, (slice, int, np.int32, np.int64, list, np.ndarray, ak.Array)):
+            if isinstance(key, (int, np.int32, np.int64)):
                 key = int(key)
             return self.__class__(
                 self._fobj,
@@ -187,10 +197,12 @@ class OfflineReader:
                 step_size=self._step_size,
                 aliases=self.aliases,
                 keys=self.keys(),
-                event_ctor=self._event_ctor
+                event_ctor=self._event_ctor,
             )
 
-        if isinstance(key, str) and key.startswith("n_"):  # group counts, for e.g. n_events, n_hits etc.
+        if isinstance(key, str) and key.startswith(
+            "n_"
+        ):  # group counts, for e.g. n_events, n_hits etc.
             key = self._keyfor(key.split("n_")[1])
             arr = self._fobj[self.event_path][key].array(uproot.AsDtype(">i4"))
             return unfold_indices(arr, self._index_chain)
@@ -207,22 +219,21 @@ class OfflineReader:
                 if from_field in branch[key].keys():
                     fields.append(to_field)
             log.debug(fields)
-            out = branch[key].arrays(
-                fields, aliases=self.special_branches[key]
-            )
+            out = branch[key].arrays(fields, aliases=self.special_branches[key])
         else:
             out = branch[self.aliases.get(key, key)].array()
 
         return unfold_indices(out, self._index_chain)
 
     def __iter__(self):
-        self._iterator_index = 0
         self._events = self._event_generator()
         return self
 
     def _event_generator(self):
         events = self._fobj[self.event_path]
-        group_count_keys = set(k for k in self.keys() if k.startswith("n_"))  # special keys to make it easy to count subbranch lengths
+        group_count_keys = set(
+            k for k in self.keys() if k.startswith("n_")
+        )  # special keys to make it easy to count subbranch lengths
         log.debug("group_count_keys: %s", group_count_keys)
         keys = set(
             list(
@@ -287,13 +298,12 @@ class OfflineReader:
             return 1
         else:
             # ignore the usual index magic and access `id` directly
-            return len(self._fobj[self.event_path]["id"].array(), self._index_chain)
+            return len(unfold_indices(self._fobj[self.event_path]["id"].array(), self._index_chain))
 
     def __actual_len__(self):
         """The raw number of events without any indexing/slicing magic"""
         return len(self._fobj[self.event_path]["id"].array())
 
-
     def __repr__(self):
         length = len(self)
         actual_length = self.__actual_len__()