expmem.py

Go to the documentation of this file.

 
import ROOT
from fnmatch import fnmatchcase
from six import string_types
 
class ExpMem(object):
    """
    Wraps exported Go4 memory .root-files
    """
 
    def __init__(self, filename, rootfolder="Analysis"):
        """
        Open the TFile filename as file and
        get the TDirectory rootfolder as root
        For rootfolder, "Workspace" might be desired besides the default
        An unset rootfolder, "*", and "/" all mean the actual root of the file
        """
        if not rootfolder or rootfolder in ["*", "/"]:
            rootfolder = ""
 
        self.file = f = ROOT.TFile(filename)
        self.root = f.GetDirectory(rootfolder)
 
        if not self.root:
            self.root = self.file
 
 
    def Objects(self, folder=None):
        """
        Iterate over all objects in the opened .root-file
        folder can be a string (relative to rootfolder) or a TDirectory
        """
        if folder is None:
            folder = self.root
        elif isinstance(folder, string_types):
            folder = self.root.GetDirectory(folder)
 
        if not folder:
            return
 
        diriter = folder.GetListOfKeys().MakeIterator()
        while True:
            key = diriter.Next()
            if not key:
                break
            if key.GetClassName().startswith("TDirectory"):
                subfolder = key.ReadObj()
                for obj in self.Objects(subfolder):
                    yield obj
            else:
                yield key.ReadObj()
 
 
    def MatchingObjects(self, expr="*", folder=None):
        """
        Iterate over Objects() that match expr
        """
        for obj in self.Objects(folder):
            name = obj.GetName()
            if fnmatchcase(name, expr):
                yield obj
 
 
    def tree(self, folder=None, indent=1):
        """
        Recursively print the object/folder tree of the opened .root-file
        """
        if folder is None:
            folder = self.root
            print(folder.GetName())
 
        diriter = folder.GetListOfKeys().MakeIterator()
        while True:
            key = diriter.Next()
            if not key:
                break
            if key.GetClassName().startswith("TDirectory"):
                subfolder = key.ReadObj()
                sfname = subfolder.GetName()
                print("-"*indent, sfname)
                self.tree(subfolder, indent+1)
            else:
                print("="*indent, key.ReadObj().GetName())