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:
      pass

  # 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,ttree):
          self._tree = ttree
          self._usedBranches = dict()
          self._index = -1
          self._aliases = ttree.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 tree(self):
          return self._tree
      def index(self):
          return self._index
      def __setBranchIndicies(self):
          for branch in self._usedBranches.itervalues():
              branch.setIndex(self._index)
      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
      def __getattr__(self, name):
          return self.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):
          print "========================================"
          print "ERROR:", message
          print "========================================"
Revision:	1.7
Committed:	Fri Feb 23 18:56:55 2007 UTC (18 years, 2 months ago) by hegner
Content type:	text/x-python
Branch:	MAIN
Changes since 1.6:	+6 -2 lines
Log Message:	improved EventTree looping. branch access now via __getattr__
#	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			pass
28
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			class EventTree(object):
59			def __init__(self,ttree):
60			self._tree = ttree
61			self._usedBranches = dict()
62			self._index = -1
63	hegner	1.5	self._aliases = ttree.GetListOfAliases()
64	hegner	1.4	def branch(self,name):
65			# support for aliases
66			alias = self._tree.GetAlias(name)
67			if alias != '': name = alias
68			# access the branch in ttree
69			if name in self._usedBranches:
70			return self._usedBranches[name]
71			self._usedBranches[name]=EventBranch(self,name)
72			return self._usedBranches[name]
73	hegner	1.5	def getListOfAliases(self):
74			return self._aliases
75	hegner	1.4	def tree(self):
76			return self._tree
77			def index(self):
78			return self._index
79			def __setBranchIndicies(self):
80			for branch in self._usedBranches.itervalues():
81			branch.setIndex(self._index)
82			def __getitem__(self,key):
83			if key <0 or key > self._tree.GetEntries():
84			raise IndexError
85			self._index = key
86			self.__setBranchIndicies()
87	hegner	1.7	return self
88	hegner	1.4	def __iter__(self):
89			for entry in xrange(self._tree.GetEntries()):
90			self._index = entry
91			self.__setBranchIndicies()
92	hegner	1.7	self._tree.GetEntry(self._index,0)
93			yield self
94			def __getattr__(self, name):
95			return self.branch(name)
96
97	hegner	1.4
98			class EventBranch(object):
99			def __init__(self,parent,name):
100			self._branch = parent.tree().GetBranch(name)
101			if self._branch == None:
102			raise cmserror("Unknown branch "+name)
103			self._buffer = createBranchBuffer(self._branch)
104			self._index = parent.index()
105			self._readData = False
106			def setIndex(self,index):
107			self._index = index
108			self._readData = False
109			def __readData(self):
110			self._branch.GetEntry(self._index)
111			self._readData = True
112
113			# replace this by __getattr__ to allow branch.attr instead of branch().attr
114			def __call__(self):
115			if not self._readData:
116			self.__readData()
117			return self._buffer
118
119			class cmserror(exceptions.StandardError):
120			def __init__(self, message):
121			print "========================================"
122			print "ERROR:", message
123			print "========================================"