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