python/framework/chain.py

0001 # Copyright (C) 2014 Colin Bernet
0002 # https://github.com/cbernet/heppy/blob/master/LICENSE
0003
0004 import glob
0005 import os
0006 import pprint
0007 from ROOT import TChain, TFile, TTree, gSystem
0008
0009 def is_pfn(fn):
0010     return not (is_lfn(fn) or is_rootfn(fn))
0011
0012 def is_lfn(fn):
0013     return fn.startswith("/store")
0014
0015 def is_rootfn(fn):
0016     """
0017     To open files like root://, file:// which os.isfile won't find.
0018     """
0019     return "://" in fn
0020
0021
0022 class Chain( object ):
0023     """Wrapper to TChain, with a python iterable interface.
0024
0025     Example of use:  #TODO make that a doctest / nose?
0026        from chain import Chain
0027        the_chain = Chain('../test/test_*.root', 'test_tree')
0028        event3 = the_chain[2]
0029        print event3.var1
0030
0031        for event in the_chain:
0032            print event.var1
0033     """
0034
0035     def __init__(self, input, tree_name=None):
0036         """
0037         Create a chain.
0038
0039         Parameters:
0040           input     = either a list of files or a wildcard (e.g. 'subdir/*.root').
0041                       In the latter case all files matching the pattern will be used
0042                       to build the chain.
0043           tree_name = key of the tree in each file.
0044                       if None and if each file contains only one TTree,
0045                       this TTree is used.
0046         """
0047         self.files = input
0048         if isinstance(input, str): # input is a pattern
0049             self.files = glob.glob(input)
0050             if len(self.files)==0:
0051                 raise ValueError('no matching file name: '+input)
0052         else: # case of a list of files
0053             if False in [
0054                 ((is_pfn(fnam) and os.path.isfile(fnam)) or
0055                 is_lfn(fnam)) or is_rootfn(fnam)
0056                 for fnam in self.files]:
0057                 err = 'at least one input file does not exist\n'
0058                 err += pprint.pformat(self.files)
0059                 raise ValueError(err)
0060         if tree_name is None:
0061             tree_name = self._guessTreeName(input)
0062         self.chain = TChain(tree_name)
0063         for file in self.files:
0064             self.chain.Add(file)
0065
0066     def _guessTreeName(self, pattern):
0067         """
0068         Find the set of keys of all TTrees in all files matching pattern.
0069         If the set contains only one key
0070           Returns: the TTree key
0071         else raises ValueError.
0072         """
0073         names = []
0074         for fnam in self.files:
0075             rfile = TFile(fnam)
0076             for key in rfile.GetListOfKeys():
0077                 obj = rfile.Get(key.GetName())
0078                 if isinstance(obj, TTree):
0079                     names.append( key.GetName() )
0080         thename = set(names)
0081         if len(thename)==1:
0082             return list(thename)[0]
0083         else:
0084             err = [
0085                 'several TTree keys in {pattern}:'.format(
0086                     pattern=pattern
0087                     ),
0088                 ','.join(thename)
0089                 ]
0090             raise ValueError('\n'.join(err))
0091
0092     def __getattr__(self, attr):
0093         """
0094         All functions of the wrapped TChain are made available
0095         """
0096         return getattr(self.chain, attr)
0097
0098     def __iter__(self):
0099         return iter(self.chain)
0100
0101     def __len__(self):
0102         return int(self.chain.GetEntries())
0103
0104     def __getitem__(self, index):
0105         """
0106         Returns the event at position index.
0107         """
0108         self.chain.GetEntry(index)
0109         return self.chain
0110
0111