TChain context manager

TChain_contextmanager.py

import ROOT

class TChain(ROOT.TChain):
    """
    A contextmanager for TChain, which loads from a file directly.
    Can disable all branches except the needed branches (default: all enabled).
    """

    def __init__(self, tree_name, file_names=None, active_branches=None):
        super(TChain, self).__init__(tree_name)
        self.file_names = file_names
        if active_branches is not None:
            self.SetBranchStatus("*", False)
            for branch in active_branches:
                self.SetBranchStatus(branch, True)

    def __enter__(self):
        for file_name in self.file_names:
            self.Add(file_name)
        return self

    def __exit__(self, *args):
        pass