Contextmanagers for ROOT.TTree

TTree_contextmanagers.py

import ROOT
import contextlib2

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

    def __init__(self, tree_name, file_name, active_branches=None):
        """
        >>> import ROOT
        >>> from utils import TTree
        >>> with TTree("DYTuple/DecayTree", "../ntuple/dy_tuple_12_mu_reduced.root", ["Z0_MM"]) as tree:
        ...     tree.Draw("Z0_MM")
        """
        self.file_name = file_name
        self.tree_name = tree_name
        self.active_branches = active_branches

    def __enter__(self):
        self.file = ROOT.TFile(self.file_name)
        tree = self.file.Get(self.tree_name)
        if self.active_branches is not None:
            tree.SetBranchStatus("*", False)
            for branch in self.active_branches:
                tree.SetBranchStatus(branch, True)
        return tree

    def __exit__(self, *args):
        self.file.Close()
    

class TTreeStack(object):
    """
    A stack of TTrees, can be used to create multiple trees at once:
    
    Example:
      >>> with TTreeStack("tree", "file1.root", "file2.root") as (tree1, tree2):
      ...    tree1.Print()
      ...    tree2.Print()
    """

    def __init__(self, tree_name, *file_names, **kwargs):
        self.tree_name = tree_name
        self.file_names = file_names
        self.active_branches = kwargs.get("active_branches", None)

    def __enter__(self):
        with contextlib2.ExitStack() as stack:
            self.trees = [stack.enter_context(
                TTree(self.tree_name, fname, active_branches=self.branches)) for fname in self.file_names]
            self.close = stack.pop_all().close
        return self.trees

    def __exit__(self, exception_type, exception_value, traceback):
        self.close()

@swang373 In that case I think a TChain is preferable, because it is already designed to behave like a TTree but pulling from multiple trees in the background. The above TTreeStack is good if you need to open all these trees but want to do different stuff with them afterwards (I need to extract templates for a fit from each tree, doing some reweighting on them first).

Also note that the TTreeStack is just a shorter way to write:
with TTree("tree", "file1.root") as tree1, TTree("tree", "file2.root") as tree2, ...:

I also wrote a contextmanager for a TChain, but haven't really needed it so far (see https://gist.github.com/graipher/f4f35d792a97c65c6c458c65c3cc9295).

@graipher That use case for distinct operations on the separate trees makes sense, I was stuck thinking of operating on them uniformly. Thanks for the examples!