PythonAnalysis/python/cmstools.py

"""Python helper tools for CMS FWLite

benedikt.hegner@cern.ch

"""
import re
import ROOT
import exceptions
### define tab completion
try:
  import readline, cmscompleter
  readline.parse_and_bind('tab: complete')
except:
  print 'WARNING: Could not load tab completion'


### workaround iterator generators for ROOT classes
def all(container):

  # loop over ROOT::TTree and similar
  if hasattr(container,'GetEntries'):
    try:
      entries = container.GetEntries()
      for entry in xrange(entries):
        yield entry
    except:
        raise cmserror("Looping of %s failed" %container) 

  # loop over std::vectors and similar
  elif hasattr(container, 'size'):
    try:
      entries = container.size()
      for entry in xrange(entries):
        yield container[entry]
    except:
      pass

 # loop over containers with begin and end iterators
def loop(begin, end):
    """Convert a pair of C++ iterators into a python generator"""
    while (begin != end):
        yield begin.__deref__()  #*b
        begin.__preinc__()       #++b 

### auto branch types (Chris Jones)
def createBranchBuffer(branch):
    reColons = re.compile(r'::')
    reCloseTemplate =re.compile(r'>')
    reOpenTemplate =re.compile(r'<')
    branchType = ROOT.branchToClass(branch)
    buffer = eval ('ROOT.'+reColons.sub(".",reOpenTemplate.sub("(ROOT.",reCloseTemplate.sub(")",branchType.GetName())))+'()')
    if( branch.GetName()[-1] != '.'):
        branch.SetAddress(buffer)
    else:
        branch.SetAddress(ROOT.AddressOf(buffer))
    return buffer


class EventTree(object):
      def __init__(self,obj):
          if isinstance(obj, ROOT.TTree):
              self._tree = obj
          elif isinstance(obj, ROOT.TFile):
              self._tree = obj.Get("Events")
          elif isinstance(obj, str):
              self._tree = ROOT.TFile.Open(obj).Get("Events")
          else:
              raise cmserror("EventTree accepts only TTrees, TFiles and filenames")
          self._usedBranches = dict()
          self._index = -1
          self._aliases = self._tree.GetListOfAliases()
      def branch(self,name):
          # support for aliases
          alias = self._tree.GetAlias(name)
          if alias != '': name = alias 
          # access the branch in ttree
          if name in self._usedBranches:
              return self._usedBranches[name]
          self._usedBranches[name]=EventBranch(self,name)
          return self._usedBranches[name]
      def getListOfAliases(self):
          return self._aliases
      def index(self):
          return self._index
      def tree(self):
          return self._tree
      def __setBranchIndicies(self):
          for branch in self._usedBranches.itervalues():
              branch.setIndex(self._index)
      def __getattr__(self, name):
          return self.branch()
      def __getitem__(self,key):
          if key <0 or key > self._tree.GetEntries():
              raise IndexError
          self._index = key
          self.__setBranchIndicies()
          return self
      def __iter__(self):
          for entry in xrange(self._tree.GetEntries()):
              self._index = entry
              self.__setBranchIndicies()
              self._tree.GetEntry(self._index,0)
              yield self
              

class EventBranch(object):
    def __init__(self,parent,name):
        self._branch = parent.tree().GetBranch(name)
        if self._branch == None:
            raise cmserror("Unknown branch "+name)
        self._buffer = createBranchBuffer(self._branch)
        self._index = parent.index()
        self._readData = False
    def setIndex(self,index):
        self._index = index
        self._readData = False
    def __readData(self):
        self._branch.GetEntry(self._index)
        self._readData = True

    # replace this by __getattr__ to allow branch.attr instead of branch().attr
    def __call__(self):
        if not self._readData:
            self.__readData()
        return self._buffer


class cmserror(exceptions.StandardError):
    def __init__(self, message):
          length = len(message)+7   #7=len("ERROR: ")
          print "="*length
          print "ERROR:", message
          print "="*length
Revision:	1.9
Committed:	Fri Feb 23 22:35:45 2007 UTC (18 years, 2 months ago) by hegner
Content type:	text/x-python
Branch:	MAIN
Changes since 1.8:	+11 -15 lines
Log Message:	code cleanup
#	User	Rev	Content
1	hegner	1.1	"""Python helper tools for CMS FWLite
2
3			benedikt.hegner@cern.ch
4
5			"""
6	hegner	1.4	import re
7			import ROOT
8			import exceptions
9	hegner	1.1	### define tab completion
10			try:
11			import readline, cmscompleter
12			readline.parse_and_bind('tab: complete')
13			except:
14			print 'WARNING: Could not load tab completion'
15
16
17			### workaround iterator generators for ROOT classes
18			def all(container):
19
20			# loop over ROOT::TTree and similar
21			if hasattr(container,'GetEntries'):
22			try:
23			entries = container.GetEntries()
24			for entry in xrange(entries):
25			yield entry
26			except:
27	hegner	1.9	raise cmserror("Looping of %s failed" %container)
28	hegner	1.1
29			# loop over std::vectors and similar
30			elif hasattr(container, 'size'):
31			try:
32			entries = container.size()
33			for entry in xrange(entries):
34			yield container[entry]
35			except:
36			pass
37	hegner	1.4
38	hegner	1.6	# loop over containers with begin and end iterators
39			def loop(begin, end):
40			"""Convert a pair of C++ iterators into a python generator"""
41			while (begin != end):
42			yield begin.__deref__() #*b
43			begin.__preinc__() #++b
44	hegner	1.4
45			### auto branch types (Chris Jones)
46			def createBranchBuffer(branch):
47			reColons = re.compile(r'::')
48			reCloseTemplate =re.compile(r'>')
49			reOpenTemplate =re.compile(r'<')
50			branchType = ROOT.branchToClass(branch)
51			buffer = eval ('ROOT.'+reColons.sub(".",reOpenTemplate.sub("(ROOT.",reCloseTemplate.sub(")",branchType.GetName())))+'()')
52			if( branch.GetName()[-1] != '.'):
53			branch.SetAddress(buffer)
54			else:
55			branch.SetAddress(ROOT.AddressOf(buffer))
56			return buffer
57
58	hegner	1.9
59	hegner	1.4	class EventTree(object):
60	hegner	1.8	def __init__(self,obj):
61			if isinstance(obj, ROOT.TTree):
62			self._tree = obj
63			elif isinstance(obj, ROOT.TFile):
64			self._tree = obj.Get("Events")
65			elif isinstance(obj, str):
66			self._tree = ROOT.TFile.Open(obj).Get("Events")
67			else:
68	hegner	1.9	raise cmserror("EventTree accepts only TTrees, TFiles and filenames")
69	hegner	1.8	self._usedBranches = dict()
70			self._index = -1
71			self._aliases = self._tree.GetListOfAliases()
72	hegner	1.4	def branch(self,name):
73			# support for aliases
74			alias = self._tree.GetAlias(name)
75			if alias != '': name = alias
76			# access the branch in ttree
77			if name in self._usedBranches:
78			return self._usedBranches[name]
79			self._usedBranches[name]=EventBranch(self,name)
80			return self._usedBranches[name]
81	hegner	1.5	def getListOfAliases(self):
82			return self._aliases
83	hegner	1.9	def index(self):
84			return self._index
85	hegner	1.4	def tree(self):
86			return self._tree
87			def __setBranchIndicies(self):
88			for branch in self._usedBranches.itervalues():
89			branch.setIndex(self._index)
90	hegner	1.9	def __getattr__(self, name):
91			return self.branch()
92	hegner	1.4	def __getitem__(self,key):
93			if key <0 or key > self._tree.GetEntries():
94			raise IndexError
95			self._index = key
96			self.__setBranchIndicies()
97	hegner	1.7	return self
98	hegner	1.4	def __iter__(self):
99			for entry in xrange(self._tree.GetEntries()):
100			self._index = entry
101			self.__setBranchIndicies()
102	hegner	1.7	self._tree.GetEntry(self._index,0)
103			yield self
104
105	hegner	1.4
106			class EventBranch(object):
107			def __init__(self,parent,name):
108			self._branch = parent.tree().GetBranch(name)
109			if self._branch == None:
110			raise cmserror("Unknown branch "+name)
111			self._buffer = createBranchBuffer(self._branch)
112			self._index = parent.index()
113			self._readData = False
114			def setIndex(self,index):
115			self._index = index
116			self._readData = False
117			def __readData(self):
118			self._branch.GetEntry(self._index)
119			self._readData = True
120
121			# replace this by __getattr__ to allow branch.attr instead of branch().attr
122			def __call__(self):
123			if not self._readData:
124			self.__readData()
125			return self._buffer
126
127	hegner	1.9
128	hegner	1.4	class cmserror(exceptions.StandardError):
129			def __init__(self, message):
130	hegner	1.9	length = len(message)+7 #7=len("ERROR: ")
131			print "="*length
132	hegner	1.4	print "ERROR:", message
133	hegner	1.9	print "="*length