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,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 __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.8
Committed:	Fri Feb 23 19:40:03 2007 UTC (18 years, 2 months ago) by hegner
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	V00-01-01
Changes since 1.7:	+14 -0 lines
Log Message:	new short cut for opening event files
#	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	hegner	1.8	def __init__(self,obj):
60			if isinstance(obj, ROOT.TTree):
61			self._tree = obj
62			elif isinstance(obj, ROOT.TFile):
63			self._tree = obj.Get("Events")
64			elif isinstance(obj, str):
65			self._tree = ROOT.TFile.Open(obj).Get("Events")
66			else:
67			raise cmserror, "EventTree accepts only TTrees, TFiles and filenames"
68			self._usedBranches = dict()
69			self._index = -1
70			self._aliases = self._tree.GetListOfAliases()
71
72
73	hegner	1.4	def __init__(self,ttree):
74			self._tree = ttree
75			self._usedBranches = dict()
76			self._index = -1
77	hegner	1.5	self._aliases = ttree.GetListOfAliases()
78	hegner	1.4	def branch(self,name):
79			# support for aliases
80			alias = self._tree.GetAlias(name)
81			if alias != '': name = alias
82			# access the branch in ttree
83			if name in self._usedBranches:
84			return self._usedBranches[name]
85			self._usedBranches[name]=EventBranch(self,name)
86			return self._usedBranches[name]
87	hegner	1.5	def getListOfAliases(self):
88			return self._aliases
89	hegner	1.4	def tree(self):
90			return self._tree
91			def index(self):
92			return self._index
93			def __setBranchIndicies(self):
94			for branch in self._usedBranches.itervalues():
95			branch.setIndex(self._index)
96			def __getitem__(self,key):
97			if key <0 or key > self._tree.GetEntries():
98			raise IndexError
99			self._index = key
100			self.__setBranchIndicies()
101	hegner	1.7	return self
102	hegner	1.4	def __iter__(self):
103			for entry in xrange(self._tree.GetEntries()):
104			self._index = entry
105			self.__setBranchIndicies()
106	hegner	1.7	self._tree.GetEntry(self._index,0)
107			yield self
108			def __getattr__(self, name):
109			return self.branch(name)
110
111	hegner	1.4
112			class EventBranch(object):
113			def __init__(self,parent,name):
114			self._branch = parent.tree().GetBranch(name)
115			if self._branch == None:
116			raise cmserror("Unknown branch "+name)
117			self._buffer = createBranchBuffer(self._branch)
118			self._index = parent.index()
119			self._readData = False
120			def setIndex(self,index):
121			self._index = index
122			self._readData = False
123			def __readData(self):
124			self._branch.GetEntry(self._index)
125			self._readData = True
126
127			# replace this by __getattr__ to allow branch.attr instead of branch().attr
128			def __call__(self):
129			if not self._readData:
130			self.__readData()
131			return self._buffer
132
133			class cmserror(exceptions.StandardError):
134			def __init__(self, message):
135			print "========================================"
136			print "ERROR:", message
137			print "========================================"