Merge pull request #39 from jpata/heppy-chain-744p4

Thanks for this improvement Joosep!

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)

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'''