diff --git a/PhysicsTools/HeppyCore/python/framework/chain.py b/PhysicsTools/HeppyCore/python/framework/chain.py index 5474b42e758a0..43117bbcba2c2 100644 --- a/PhysicsTools/HeppyCore/python/framework/chain.py +++ b/PhysicsTools/HeppyCore/python/framework/chain.py @@ -6,6 +6,19 @@ import pprint from ROOT import TChain, TFile, TTree, gSystem +def is_pfn(fn): + return not (is_lfn(fn) or is_rootfn(fn)) + +def is_lfn(fn): + return fn.startswith("/store") + +def is_rootfn(fn): + """ + To open files like root://, file:// which os.isfile won't find. + """ + return "://" in fn + + class Chain( object ): """Wrapper to TChain, with a python iterable interface. @@ -37,7 +50,10 @@ def __init__(self, input, tree_name=None): if len(self.files)==0: raise ValueError('no matching file name: '+input) else: # case of a list of files - if False in [ os.path.isfile(fnam) for fnam in self.files ]: + if False in [ + ((is_pfn(fnam) and os.path.isfile(fnam)) or + is_lfn(fnam)) or is_rootfn(fnam) + for fnam in self.files]: err = 'at least one input file does not exist\n' err += pprint.pformat(self.files) raise ValueError(err) diff --git a/PhysicsTools/HeppyCore/python/framework/chain_test.py b/PhysicsTools/HeppyCore/python/framework/chain_test.py index e124488d18bda..ccc308f229fea 100644 --- a/PhysicsTools/HeppyCore/python/framework/chain_test.py +++ b/PhysicsTools/HeppyCore/python/framework/chain_test.py @@ -35,6 +35,11 @@ def test_load_2(self): chain = Chain(testfname.replace('.root', '*.root'), 'test_tree') self.assertEqual(len(chain), 200) os.remove(tmpfile) + + def test_load_3(self): + '''Test LFN/root-fn loading''' + chain = Chain(["root://{0}".format(os.path.abspath(testfname))], 'test_tree') + self.assertEqual(len(chain), 100) def test_iterate(self): '''Test iteration'''