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'


# for adding iterators at runtime
import iterators


### 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())))+'()')
    buffer = ROOT.MakeRootClass(branchType.GetName()) ()
    if( branch.GetName()[-1] != '.') and (branch.GetName()!="EventAuxiliary"):
        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 cppCode(self, name):
          """C++ code for accessing the product inside the full framework"""
          alias = self._tree.GetAlias(name)
          if alias != '': name = alias
          tmpBranch = self._tree.GetBranch(name)
          typeString = ROOT.branchToClass(tmpBranch).GetName()
          if "edm::Wrapper" in typeString:
              typeString = typeString.replace("<edm::Wrapper","")
              typeString = typeString.rstrip(">")
          nameParts = name.split("_")
          if nameParts[2] == "":
              cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", dummy);'\
                        %(typeString, nameParts[1])
          else:
              cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", "%s", dummy);'\
                        %(typeString, nameParts[1], nameParts[2])
          return cppCode
      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(name)
      def __getitem__(self,key):
          if key <0 or key > self._tree.GetEntries():
              raise IndexError
          self._index = key
          self.__setBranchIndicies()
          self._tree.GetEntry(self._index,0)
          return Event(self)
      def __iter__(self):
          # flushing/initializing the root buffers 
          entry = 0
          self._index = entry
          self.__setBranchIndicies()
          self._tree.GetEntry(self._index,0)
          # the real loop
          for entry in xrange(self._tree.GetEntries()):
              self._index = entry
              self.__setBranchIndicies()
              self._tree.GetEntry(self._index,0)
              yield Event(self)    # TODO: don't return a new object but update the old one 
              

class Event(object):
    def __init__(self, eventTree):
        self._eventTree = eventTree

    def getProduct(self, name):
        return iterators.addIterator(self._eventTree.branch(name)())

    def __getattr__(self, name):
        return iterators.addIterator(self._eventTree.branch(name)())


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.17
Committed:	Mon Nov 3 16:11:46 2008 UTC (16 years, 6 months ago) by hegner
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	V00-02-02
Changes since 1.16:	+1 -1 lines
Log Message:	fix EventAux problem
#	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	hegner	1.10	# for adding iterators at runtime
18			import iterators
19
20
21	hegner	1.1	### workaround iterator generators for ROOT classes
22			def all(container):
23
24			# loop over ROOT::TTree and similar
25			if hasattr(container,'GetEntries'):
26			try:
27			entries = container.GetEntries()
28			for entry in xrange(entries):
29			yield entry
30			except:
31	hegner	1.9	raise cmserror("Looping of %s failed" %container)
32	hegner	1.1
33			# loop over std::vectors and similar
34			elif hasattr(container, 'size'):
35			try:
36			entries = container.size()
37			for entry in xrange(entries):
38			yield container[entry]
39			except:
40			pass
41	hegner	1.4
42	hegner	1.6	# loop over containers with begin and end iterators
43			def loop(begin, end):
44			"""Convert a pair of C++ iterators into a python generator"""
45			while (begin != end):
46			yield begin.__deref__() #*b
47			begin.__preinc__() #++b
48	hegner	1.4
49			### auto branch types (Chris Jones)
50			def createBranchBuffer(branch):
51			reColons = re.compile(r'::')
52			reCloseTemplate =re.compile(r'>')
53			reOpenTemplate =re.compile(r'<')
54			branchType = ROOT.branchToClass(branch)
55	hegner	1.13	#buffer = eval ('ROOT.'+reColons.sub(".",reOpenTemplate.sub("(ROOT.",reCloseTemplate.sub(")",branchType.GetName())))+'()')
56			buffer = ROOT.MakeRootClass(branchType.GetName()) ()
57	hegner	1.17	if( branch.GetName()[-1] != '.') and (branch.GetName()!="EventAuxiliary"):
58	hegner	1.4	branch.SetAddress(buffer)
59			else:
60			branch.SetAddress(ROOT.AddressOf(buffer))
61			return buffer
62
63	hegner	1.9
64	hegner	1.4	class EventTree(object):
65	hegner	1.8	def __init__(self,obj):
66			if isinstance(obj, ROOT.TTree):
67			self._tree = obj
68			elif isinstance(obj, ROOT.TFile):
69			self._tree = obj.Get("Events")
70			elif isinstance(obj, str):
71			self._tree = ROOT.TFile.Open(obj).Get("Events")
72			else:
73	hegner	1.9	raise cmserror("EventTree accepts only TTrees, TFiles and filenames")
74	hegner	1.8	self._usedBranches = dict()
75			self._index = -1
76			self._aliases = self._tree.GetListOfAliases()
77	hegner	1.4	def branch(self,name):
78			# support for aliases
79			alias = self._tree.GetAlias(name)
80			if alias != '': name = alias
81			# access the branch in ttree
82			if name in self._usedBranches:
83			return self._usedBranches[name]
84			self._usedBranches[name]=EventBranch(self,name)
85			return self._usedBranches[name]
86	hegner	1.11	def cppCode(self, name):
87			"""C++ code for accessing the product inside the full framework"""
88			alias = self._tree.GetAlias(name)
89			if alias != '': name = alias
90			tmpBranch = self._tree.GetBranch(name)
91			typeString = ROOT.branchToClass(tmpBranch).GetName()
92	hegner	1.12	if "edm::Wrapper" in typeString:
93			typeString = typeString.replace("<edm::Wrapper","")
94			typeString = typeString.rstrip(">")
95	hegner	1.11	nameParts = name.split("_")
96			if nameParts[2] == "":
97			cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", dummy);'\
98			%(typeString, nameParts[1])
99			else:
100			cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", "%s", dummy);'\
101			%(typeString, nameParts[1], nameParts[2])
102			return cppCode
103	hegner	1.5	def getListOfAliases(self):
104			return self._aliases
105	hegner	1.9	def index(self):
106			return self._index
107	hegner	1.4	def tree(self):
108			return self._tree
109			def __setBranchIndicies(self):
110			for branch in self._usedBranches.itervalues():
111			branch.setIndex(self._index)
112	hegner	1.9	def __getattr__(self, name):
113	hegner	1.15	return self.branch(name)
114	hegner	1.4	def __getitem__(self,key):
115			if key <0 or key > self._tree.GetEntries():
116			raise IndexError
117			self._index = key
118			self.__setBranchIndicies()
119	hegner	1.16	self._tree.GetEntry(self._index,0)
120			return Event(self)
121	hegner	1.4	def __iter__(self):
122	hegner	1.14	# flushing/initializing the root buffers
123			entry = 0
124			self._index = entry
125			self.__setBranchIndicies()
126			self._tree.GetEntry(self._index,0)
127			# the real loop
128	hegner	1.4	for entry in xrange(self._tree.GetEntries()):
129			self._index = entry
130			self.__setBranchIndicies()
131	hegner	1.7	self._tree.GetEntry(self._index,0)
132	hegner	1.10	yield Event(self) # TODO: don't return a new object but update the old one
133	hegner	1.7
134	hegner	1.4
135	hegner	1.10	class Event(object):
136			def __init__(self, eventTree):
137			self._eventTree = eventTree
138
139	hegner	1.13	def getProduct(self, name):
140	hegner	1.10	return iterators.addIterator(self._eventTree.branch(name)())
141
142			def __getattr__(self, name):
143			return iterators.addIterator(self._eventTree.branch(name)())
144
145
146	hegner	1.4	class EventBranch(object):
147			def __init__(self,parent,name):
148			self._branch = parent.tree().GetBranch(name)
149			if self._branch == None:
150			raise cmserror("Unknown branch "+name)
151			self._buffer = createBranchBuffer(self._branch)
152			self._index = parent.index()
153			self._readData = False
154			def setIndex(self,index):
155			self._index = index
156			self._readData = False
157			def __readData(self):
158			self._branch.GetEntry(self._index)
159			self._readData = True
160
161			# replace this by __getattr__ to allow branch.attr instead of branch().attr
162			def __call__(self):
163			if not self._readData:
164			self.__readData()
165			return self._buffer
166
167	hegner	1.9
168	hegner	1.4	class cmserror(exceptions.StandardError):
169			def __init__(self, message):
170	hegner	1.9	length = len(message)+7 #7=len("ERROR: ")
171			print "="*length
172	hegner	1.4	print "ERROR:", message
173	hegner	1.9	print "="*length