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] != '.'):
        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()
      def __getitem__(self,key):
          if key <0 or key > self._tree.GetEntries():
              raise IndexError
          self._index = key
          self.__setBranchIndicies()
          return 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.14
Committed:	Wed Jun 13 13:02:30 2007 UTC (17 years, 10 months ago) by hegner
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CMSSW_1_6_12, CMSSW_2_0_8, CMSSW_1_6_12_pre2, CMSSW_1_6_12_pre1, CMSSW_2_0_7, CMSSW_2_0_6, CMSSW_2_0_5, CMSSW_2_0_4, CMSSW_2_0_3, CMSSW_2_0_2, CMSSW_2_0_1, CMSSW_2_0_0, CMSSW_2_1_0_pre1, CMSSW_1_8_4, CMSSW_2_0_0_pre9, CMSSW_2_0_0_pre8, CMSSW_1_8_3, CMSSW_1_6_11, CMSSW_2_0_0_pre7, CMSSW_1_8_2, CMSSW_1_7_7, CMSSW_2_0_0_pre6, CMSSW_1_6_11_pre1, CMSSW_2_0_0_pre5, CMSSW_1_8_1, CMSSW_2_0_0_pre4, CMSSW_1_8_0, CMSSW_1_7_6, CMSSW_2_0_0_pre3, CMSSW_1_6_10, CMSSW_1_8_0_pre10, CMSSW_1_6_10_pre2, CMSSW_1_6_10_pre1, CMSSW_1_8_0_pre9, CMSSW_2_0_0_pre2, CMSSW_2_0_0_pre1, CMSSW_1_6_9, CMSSW_1_8_0_pre8, CMSSW_1_6_9_pre2, CMSSW_1_8_0_pre7, CMSSW_1_6_9_pre1, CMSSW_1_7_5, CMSSW_1_8_0_pre6, CMSSW_1_8_0_pre5, CMSSW_1_7_4, CMSSW_1_7_3, CMSSW_1_8_0_pre4, CMSSW_1_6_8, CMSSW_1_6_8_pre2, CMSSW_1_8_0_pre3a, CMSSW_1_7_2, CMSSW_1_8_0_pre2, CMSSW_1_6_8_pre1, CMSSW_1_7_1, CMSSW_1_8_0_pre1, CMSSW_1_7_0, CMSSW_1_7_0_pre13, CMSSW_1_7_0_pre12, CMSSW_1_7_0_pre11, CMSSW_1_7_0_pre10, CMSSW_1_7_0_pre9, CMSSW_1_7_0_pre8, CMSSW_1_7_0_pre7, CMSSW_1_6_7, CMSSW_1_7_0_pre6, CMSSW_1_6_6, CMSSW_1_6_5, CMSSW_1_6_5_pre1, CMSSW_1_6_4, CMSSW_1_7_0_pre5, CMSSW_1_6_3, CMSSW_1_6_1, CMSSW_1_7_0_pre4, CMSSW_1_7_0_pre3, CMSSW_1_7_0_pre2, CMSSW_1_6_0, CMSSW_1_6_0_pre14, CMSSW_1_7_0_pre1, CMSSW_1_6_0_DAQ3, CMSSW_1_6_0_pre13, CMSSW_1_6_0_pre12, CMSSW_1_6_0_pre11, CMSSW_1_6_0_pre10, CMSSW_1_6_0_pre9, CMSSW_1_6_0_pre8, CMSSW_1_6_0_pre7, CMSSW_1_6_0_pre6, CMSSW_1_6_0_DAQ1, CMSSW_1_6_0_pre5, V00-01-09, CMSSW_1_6_0_pre4, CMSSW_1_6_0_pre3, CMSSW_1_6_0_pre2, V00-01-08
Changes since 1.13:	+6 -0 lines
Log Message:	bugfix for buffer problem (reported by dave)
#	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.4	if( branch.GetName()[-1] != '.'):
58			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			return self.branch()
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.7	return self
120	hegner	1.4	def __iter__(self):
121	hegner	1.14	# flushing/initializing the root buffers
122			entry = 0
123			self._index = entry
124			self.__setBranchIndicies()
125			self._tree.GetEntry(self._index,0)
126			# the real loop
127	hegner	1.4	for entry in xrange(self._tree.GetEntries()):
128			self._index = entry
129			self.__setBranchIndicies()
130	hegner	1.7	self._tree.GetEntry(self._index,0)
131	hegner	1.10	yield Event(self) # TODO: don't return a new object but update the old one
132	hegner	1.7
133	hegner	1.4
134	hegner	1.10	class Event(object):
135			def __init__(self, eventTree):
136			self._eventTree = eventTree
137
138	hegner	1.13	def getProduct(self, name):
139	hegner	1.10	return iterators.addIterator(self._eventTree.branch(name)())
140
141			def __getattr__(self, name):
142			return iterators.addIterator(self._eventTree.branch(name)())
143
144
145	hegner	1.4	class EventBranch(object):
146			def __init__(self,parent,name):
147			self._branch = parent.tree().GetBranch(name)
148			if self._branch == None:
149			raise cmserror("Unknown branch "+name)
150			self._buffer = createBranchBuffer(self._branch)
151			self._index = parent.index()
152			self._readData = False
153			def setIndex(self,index):
154			self._index = index
155			self._readData = False
156			def __readData(self):
157			self._branch.GetEntry(self._index)
158			self._readData = True
159
160			# replace this by __getattr__ to allow branch.attr instead of branch().attr
161			def __call__(self):
162			if not self._readData:
163			self.__readData()
164			return self._buffer
165
166	hegner	1.9
167	hegner	1.4	class cmserror(exceptions.StandardError):
168			def __init__(self, message):
169	hegner	1.9	length = len(message)+7 #7=len("ERROR: ")
170			print "="*length
171	hegner	1.4	print "ERROR:", message
172	hegner	1.9	print "="*length