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()

@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!