util.py

Go to the documentation of this file.

#!/usr/bin/env python2
 
import re
 
import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
 
def unpack_tdirectory(tdir):
    if tdir.GetListOfKeys():
        for key in tdir.GetListOfKeys():
            yield key.ReadObj()
    else:
        for obj in tdir.GetList():
            yield obj
 
def update_tcanvases(objects=None):
    if objects is None:
        objects = ROOT.gROOT.GetListOfCanvases()
 
    for obj in objects:
        if isinstance(obj, ROOT.TPad):
            obj.Modified()
            obj.Update()
            update_tcanvases(obj.GetListOfPrimitives())
 
 
def increment_name(name):
    res = re.search('[0-9]+$', name)
    if res:
        prefix = name[:-len(res.group())]
        number = int(res.group()) + 1
        return prefix + str(number)
    else:
        return name + '_1'
 
 
class PreserveGDir(object):
 
    def __init__(self, directory = None):
        self.directory = directory
 
    def __enter__(self):
        self.gdir = ROOT.gDirectory
        if self.directory is not None:
            self.directory.cd()
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.gdir.cd()
 
 
class TKeyDict(dict):
    def __getitem__(self, key):
        output = super(TKeyDict,self).__getitem__(key)
        if isinstance(output, ROOT.TKey):
            output = output.ReadObj()
            if (isinstance(output, ROOT.TH2) and
                not isinstance(output, ROOT.GH2Base)):
                output = ROOT.GH2D(output)
            self[key] = output
 
        return output
 
    def is_tkey(self, key):
        output = super(TKeyDict,self).__getitem__(key)
        try:
            return isinstance(output, ROOT.TKey)
        except AttributeError:
            return False