Skip to content
Snippets Groups Projects
Commit 3a913f81 authored by Tamas Gal's avatar Tamas Gal :speech_balloon:
Browse files

Cleanup OfflineReader interface

parent 7822134f
No related branches found
No related tags found
1 merge request!21Resolve "Reduce the amount of uproot.open (to one)"
...@@ -201,7 +201,7 @@ class OfflineKeys: ...@@ -201,7 +201,7 @@ class OfflineKeys:
class OfflineReader: class OfflineReader:
"""reader for offline ROOT files""" """reader for offline ROOT files"""
def __init__(self, file_path=None): def __init__(self, file_path=None, fobj=None, data=None):
""" OfflineReader class is an offline ROOT file wrapper """ OfflineReader class is an offline ROOT file wrapper
Parameters Parameters
...@@ -209,34 +209,37 @@ class OfflineReader: ...@@ -209,34 +209,37 @@ class OfflineReader:
file_path : path-like object file_path : path-like object
Path to the file of interest. It can be a str or any python Path to the file of interest. It can be a str or any python
path-like object that points to the file. path-like object that points to the file.
""" """
self._file_path = file_path self._file_path = file_path
if file_path is not None: if file_path is not None:
self._tree = uproot.open(self._file_path)[MAIN_TREE_NAME] self._fobj = uproot.open(self._file_path)
self._tree = self._fobj[MAIN_TREE_NAME]
self._data = self._tree.lazyarrays( self._data = self._tree.lazyarrays(
basketcache=uproot.cache.ThreadSafeArrayCache( basketcache=uproot.cache.ThreadSafeArrayCache(
BASKET_CACHE_SIZE)) BASKET_CACHE_SIZE))
else:
self._fobj = fobj
self._tree = self._fobj[MAIN_TREE_NAME]
self._data = data
@classmethod @classmethod
def from_tree(cls, tree, data): def from_index(cls, source, index):
instance = cls() instance = cls(fobj=source._fobj, data=source._data[index])
instance._tree = tree
instance._data = data
return instance return instance
def __getitem__(self, item): def __getitem__(self, index):
return OfflineReader.from_tree(tree=self._tree, data=self._data[item]) return OfflineReader.from_index(source=self, index=index)
def __len__(self): def __len__(self):
return len(self._data) return len(self._data)
@cached_property @cached_property
def header(self): def header(self):
fobj = uproot.open(self._file_path) if 'Head' in self._fobj:
if 'Head' in fobj:
header = {} header = {}
for n, x in fobj['Head']._map_3c_string_2c_string_3e_.items(): for n, x in self._fobj['Head']._map_3c_string_2c_string_3e_.items():
header[n.decode("utf-8")] = x.decode("utf-8").strip() header[n.decode("utf-8")] = x.decode("utf-8").strip()
return header return header
else: else:
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment