class VariableProducer(Module): def __init__(self, name, booker, modules): Module.__init__(self, name, booker) self._modules = [(n, m() if type(m) == types.FunctionType else m) for (n, m) in modules] def init(self, tree): for n, m in self._modules: if hasattr(m, 'init'): m.init(tree) def beginJob(self): self.t = PyTree(self.book("TTree", "t", "t")) self.branches = {} for name, mod in self._modules: print name print mod.listBranches() if hasattr(mod, 'setOutputTree'): mod.setOutputTree(self.t) for B in mod.listBranches(): # don't add the same branch twice if B in self.branches: print "Will not add branch %s twice" % (B, ) continue self.branches[B] = True if type(B) == tuple: if len(B) == 2: self.t.branch(B[0], B[1]) elif len(B) == 4: self.t.branch(B[0], B[1], n=B[2], lenVar=B[3]) elif len(B) == 3: self.t.branch(B[0], B[1], n=B[2], lenVar=None) else: self.t.branch(B, "F") def endJob(self): for n, m in self._modules: if hasattr(m, 'endJob'): m.endJob() def analyze(self, event): for name, mod in self._modules: keyvals = mod(event) for B, V in keyvals.iteritems(): setattr(self.t, B, V) setattr(event, B, V) self.t.fill()
def beginJob(self): self.t = PyTree(self.book("TTree", "t", "t")) self.branches = {} for name, mod in self._modules: print name print mod.listBranches() if hasattr(mod, 'setOutputTree'): mod.setOutputTree(self.t) for B in mod.listBranches(): # don't add the same branch twice if B in self.branches: print "Will not add branch %s twice" % (B, ) continue self.branches[B] = True if type(B) == tuple: if len(B) == 2: self.t.branch(B[0], B[1]) elif len(B) == 4: self.t.branch(B[0], B[1], n=B[2], lenVar=B[3]) elif len(B) == 3: self.t.branch(B[0], B[1], n=B[2], lenVar=None) else: self.t.branch(B, "F")