Use subclassed Branch

km3io/offline.py

@@ -72,6 +72,111 @@ SUBBRANCH_MAPS = [
 ]
 
 
+class OfflineBranch(Branch):
+    @cached_property
+    def usr(self):
+        return Usr(self._mapper, self._branch, index=self._index)
+
+
+class Usr:
+    """Helper class to access AAObject `usr` stuff"""
+    def __init__(self, mapper, branch, index=None):
+        self._mapper = mapper
+        self._name = mapper.name
+        self._index = index
+        self._branch = branch
+        self._usr_names = []
+        self._usr_idx_lookup = {}
+
+        self._usr_key = 'usr' if mapper.flat else mapper.key + '.usr'
+
+        self._initialise()
+
+    def _initialise(self):
+        try:
+            self._branch[self._usr_key]
+            # This will raise a KeyError in old aanet files
+            # which has a different strucuter and key (usr_data)
+            # We do not support those (yet)
+        except (KeyError, IndexError):
+            print("The `usr` fields could not be parsed for the '{}' branch.".
+                  format(self._name))
+            return
+
+        if self._mapper.flat:
+            self._initialise_flat()
+
+    def _initialise_flat(self):
+        # Here, we assume that every event has the same names in the same order
+        # to massively increase the performance. This needs triple check if
+        # it's always the case.
+        self._usr_names = [
+            n.decode("utf-8")
+            for n in self._branch[self._usr_key + '_names'].lazyarray()[0]
+        ]
+        self._usr_idx_lookup = {
+            name: index
+            for index, name in enumerate(self._usr_names)
+        }
+
+        data = self._branch[self._usr_key].lazyarray()
+
+        if self._index is not None:
+            data = data[self._index]
+
+        self._usr_data = data
+
+        for name in self._usr_names:
+            setattr(self, name, self[name])
+
+    # def _initialise_nested(self):
+    #     self._usr_names = [
+    #         n.decode("utf-8") for n in self.branch['usr_names'].lazyarray(
+    #             # TODO this will be fixed soon in uproot,
+    #             # see https://github.com/scikit-hep/uproot/issues/465
+    #             uproot.asgenobj(
+    #                 uproot.SimpleArray(uproot.STLVector(uproot.STLString())),
+    #                 self.branch['usr_names']._context, 6),
+    #             basketcache=BASKET_CACHE)[0]
+    #     ]
+
+    def __getitem__(self, item):
+        if self._mapper.flat:
+            return self.__getitem_flat__(item)
+        return self.__getitem_nested__(item)
+
+    def __getitem_flat__(self, item):
+        if self._index is not None:
+            return self._usr_data[self._index][:, self._usr_idx_lookup[item]]
+        else:
+            return self._usr_data[:, self._usr_idx_lookup[item]]
+
+    def __getitem_nested__(self, item):
+        data = self._branch[self._usr_key + '_names'].lazyarray(
+            # TODO this will be fixed soon in uproot,
+            # see https://github.com/scikit-hep/uproot/issues/465
+            uproot.asgenobj(
+                uproot.SimpleArray(uproot.STLVector(uproot.STLString())),
+                self._branch[self._usr_key + '_names']._context, 6),
+            basketcache=BASKET_CACHE)
+        if self._index is None:
+            return data
+        else:
+            return data[self._index]
+
+    def keys(self):
+        return self._usr_names
+
+    def __str__(self):
+        entries = []
+        for name in self.keys():
+            entries.append("{}: {}".format(name, self[name]))
+        return '\n'.join(entries)
+
+    def __repr__(self):
+        return "<{}[{}]>".format(self.__class__.__name__, self._name)
+
+
 class OfflineReader:
     """reader for offline ROOT files"""
     def __init__(self, file_path=None):
@@ -90,7 +195,7 @@ class OfflineReader:
     @cached_property
     def events(self):
         """The `E` branch, containing all offline events."""
-        return Branch(self._tree,
+        return OfflineBranch(self._tree,
                       mapper=EVENTS_MAP,
                       subbranchmaps=SUBBRANCH_MAPS)