EWP-BplusToKstMuMu-AngAna/Code/Scripts/Python Scripts/rootio.py

import os
import re
import uproot


class RootIO:
    def __init__(self, path_to_file: str):
        if not os.path.exists(path = path_to_file):
            raise ValueError("Path does not exists!")
        self.path_to_file = path_to_file


    def get_information(self, regex: str = None) -> None:
        with uproot.open(self.path_to_file) as file:
            print(f"File classnames:\n{file.classnames()}")
            for classname, classtype in file.classnames().items():
                print(f"Classname:\n{classname}")
                print(f"Classtype:\n{classtype}")
                print(f"Content of the class:")
                for name, branch in file[classname].items():
                    print(f"{name}: {branch}")
                if regex != None:
                    print(f"Regex search results:")
                    attributes = file[classname].keys()
                    for attribute in attributes:
                        result = self.__find_pattern_in_text(regex = regex, text = attribute)
                        if result != "":
                            print(f"{result}")


    def __find_pattern_in_text(self, regex: str, text: str) -> str:
        search_result = re.search(pattern = regex, string = text)
        result        = ""
        if search_result != None:
            result = search_result.group(0)

        return result