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())))+'()')
    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()
          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):
          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(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.11
Committed:	Fri Mar 2 02:35:35 2007 UTC (18 years, 2 months ago) by hegner
Content type:	text/x-python
Branch:	MAIN
Changes since 1.10:	+14 -0 lines
Log Message:	new EventTree.cppCode(name). It creates C++ code snippets for FW Modules.
#	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			buffer = eval ('ROOT.'+reColons.sub(".",reOpenTemplate.sub("(ROOT.",reCloseTemplate.sub(")",branchType.GetName())))+'()')
56			if( branch.GetName()[-1] != '.'):
57			branch.SetAddress(buffer)
58			else:
59			branch.SetAddress(ROOT.AddressOf(buffer))
60			return buffer
61
62	hegner	1.9
63	hegner	1.4	class EventTree(object):
64	hegner	1.8	def __init__(self,obj):
65			if isinstance(obj, ROOT.TTree):
66			self._tree = obj
67			elif isinstance(obj, ROOT.TFile):
68			self._tree = obj.Get("Events")
69			elif isinstance(obj, str):
70			self._tree = ROOT.TFile.Open(obj).Get("Events")
71			else:
72	hegner	1.9	raise cmserror("EventTree accepts only TTrees, TFiles and filenames")
73	hegner	1.8	self._usedBranches = dict()
74			self._index = -1
75			self._aliases = self._tree.GetListOfAliases()
76	hegner	1.4	def branch(self,name):
77			# support for aliases
78			alias = self._tree.GetAlias(name)
79			if alias != '': name = alias
80			# access the branch in ttree
81			if name in self._usedBranches:
82			return self._usedBranches[name]
83			self._usedBranches[name]=EventBranch(self,name)
84			return self._usedBranches[name]
85	hegner	1.11	def cppCode(self, name):
86			"""C++ code for accessing the product inside the full framework"""
87			alias = self._tree.GetAlias(name)
88			if alias != '': name = alias
89			tmpBranch = self._tree.GetBranch(name)
90			typeString = ROOT.branchToClass(tmpBranch).GetName()
91			nameParts = name.split("_")
92			if nameParts[2] == "":
93			cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", dummy);'\
94			%(typeString, nameParts[1])
95			else:
96			cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", "%s", dummy);'\
97			%(typeString, nameParts[1], nameParts[2])
98			return cppCode
99	hegner	1.5	def getListOfAliases(self):
100			return self._aliases
101	hegner	1.9	def index(self):
102			return self._index
103	hegner	1.4	def tree(self):
104			return self._tree
105			def __setBranchIndicies(self):
106			for branch in self._usedBranches.itervalues():
107			branch.setIndex(self._index)
108	hegner	1.9	def __getattr__(self, name):
109			return self.branch()
110	hegner	1.4	def __getitem__(self,key):
111			if key <0 or key > self._tree.GetEntries():
112			raise IndexError
113			self._index = key
114			self.__setBranchIndicies()
115	hegner	1.7	return self
116	hegner	1.4	def __iter__(self):
117			for entry in xrange(self._tree.GetEntries()):
118			self._index = entry
119			self.__setBranchIndicies()
120	hegner	1.7	self._tree.GetEntry(self._index,0)
121	hegner	1.10	yield Event(self) # TODO: don't return a new object but update the old one
122	hegner	1.7
123	hegner	1.4
124	hegner	1.10	class Event(object):
125			def __init__(self, eventTree):
126			self._eventTree = eventTree
127
128			def getProduct(name):
129			return iterators.addIterator(self._eventTree.branch(name)())
130
131			def __getattr__(self, name):
132			return iterators.addIterator(self._eventTree.branch(name)())
133
134
135	hegner	1.4	class EventBranch(object):
136			def __init__(self,parent,name):
137			self._branch = parent.tree().GetBranch(name)
138			if self._branch == None:
139			raise cmserror("Unknown branch "+name)
140			self._buffer = createBranchBuffer(self._branch)
141			self._index = parent.index()
142			self._readData = False
143			def setIndex(self,index):
144			self._index = index
145			self._readData = False
146			def __readData(self):
147			self._branch.GetEntry(self._index)
148			self._readData = True
149
150			# replace this by __getattr__ to allow branch.attr instead of branch().attr
151			def __call__(self):
152			if not self._readData:
153			self.__readData()
154			return self._buffer
155
156	hegner	1.9
157	hegner	1.4	class cmserror(exceptions.StandardError):
158			def __init__(self, message):
159	hegner	1.9	length = len(message)+7 #7=len("ERROR: ")
160			print "="*length
161	hegner	1.4	print "ERROR:", message
162	hegner	1.9	print "="*length