21 lines
452 B
Python
21 lines
452 B
Python
|
import numpy as np
|
||
|
|
import uproot
|
||
|
|
import typing
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class FileManager:
|
||
|
|
file: str
|
||
|
|
tree: str
|
||
|
|
branches: list[str]
|
||
|
|
|
||
|
|
def extract_data(self) -> list[np.ndarray]:
|
||
|
|
with uproot.open(self.file) as file:
|
||
|
|
file_tree = file[self.tree]
|
||
|
|
data = []
|
||
|
|
for branch in self.branches:
|
||
|
|
data.append(file_tree[branch].array(library = "np"))
|
||
|
|
|
||
|
|
return data
|