Merge pull request #837 from chrispap95/list_of_files

feat: Pass list of files to NanoEventsFactory
CoffeaTeam · Jun 14, 2023 · 7aa88ad · 7aa88ad
2 parents b8b7f20 + 34f719e
commit 7aa88ad
Show file tree

Hide file tree

Showing 8 changed files with 16 additions and 17 deletions.
diff --git a/src/coffea/nanoevents/factory.py b/src/coffea/nanoevents/factory.py
@@ -247,10 +247,11 @@ def from_root(
 
         Parameters
         ----------
-            file : str or uproot.reading.ReadOnlyDirectory
-                The filename or already opened file using e.g. ``uproot.open()``
+            file : a string or dict input to ``uproot.open()`` or ``uproot.dask()`` or a ``uproot.reading.ReadOnlyDirectory``
+                The filename or dict of filenames including the treepath (as it would be passed directly to ``uproot.open()``
+                or ``uproot.dask()``) already opened file using e.g. ``uproot.open()``.
             treepath : str, optional
-                Name of the tree to read in the file
+                Name of the tree to read in the file. Used only if ``file`` is a ``uproot.reading.ReadOnlyDirectory``.
             entry_start : int, optional
                 Start at this entry offset in the tree (default 0)
             entry_stop : int, optional
@@ -329,7 +330,7 @@ def from_root(
             else:
                 opener = partial(
                     uproot.dask,
-                    {file: treepath},
+                    file,
                     full_paths=True,
                     open_files=False,
                     ak_add_doc=True,
@@ -342,17 +343,15 @@ def from_root(
                 f"{schemaclass} is not dask capable despite allowing dask, generating non-dask nanoevents"
             )
 
-        if isinstance(file, str):
-            tree = uproot.open({file: None}, **uproot_options)[treepath]
-        elif isinstance(file, uproot.reading.ReadOnlyDirectory):
+        if isinstance(file, uproot.reading.ReadOnlyDirectory):
             tree = file[treepath]
         elif "<class 'uproot.rootio.ROOTDirectory'>" == str(type(file)):
             raise RuntimeError(
                 "The file instance (%r) is an uproot3 type, but this module is only compatible with uproot4 or higher"
                 % file
             )
         else:
-            raise TypeError("Invalid file type (%s)" % (str(type(file))))
+            tree = uproot.open(file, **uproot_options)
 
         if entry_start is None or entry_start < 0:
             entry_start = 0

diff --git a/tests/test_fix823.py b/tests/test_fix823.py
@@ -3,7 +3,9 @@
 
 def test_explicit_delete_after_assign():
     events = NanoEventsFactory.from_root(
-        "https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root",
+        {
+            "https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"
+        },
         metadata={"dataset": "nano_dy"},
         schemaclass=NanoAODSchema,
         permit_dask=True,

diff --git a/tests/test_jetmet_tools.py b/tests/test_jetmet_tools.py
@@ -673,8 +673,7 @@ def test_corrected_jets_factory():
     from coffea.nanoevents import NanoEventsFactory
 
     events = NanoEventsFactory.from_root(
-        os.path.abspath("tests/samples/nano_dy.root"),
-        treepath="Events",
+        {os.path.abspath("tests/samples/nano_dy.root"): "Events"},
         metadata={},
         permit_dask=True,
     ).events()

diff --git a/tests/test_nanoevents_delphes.py b/tests/test_nanoevents_delphes.py
@@ -10,7 +10,7 @@
 def _events():
     path = os.path.abspath("tests/samples/delphes.root")
     factory = NanoEventsFactory.from_root(
-        path, treepath="Delphes", schemaclass=DelphesSchema, permit_dask=True
+        {path: "Delphes"}, schemaclass=DelphesSchema, permit_dask=True
     )
     return factory.events()
 

diff --git a/tests/test_nanoevents_pdune.py b/tests/test_nanoevents_pdune.py
@@ -26,7 +26,7 @@
 def events():
     path = os.path.abspath("tests/samples/pduneana.root")
     factory = NanoEventsFactory.from_root(
-        path, treepath="pduneana/beamana", schemaclass=PDUNESchema
+        {path: "pduneana/beamana"}, schemaclass=PDUNESchema
     )
     return factory.events()
 

diff --git a/tests/test_nanoevents_pfnano.py b/tests/test_nanoevents_pfnano.py
@@ -9,7 +9,7 @@
 def events():
     path = os.path.abspath("tests/samples/pfnano.root")
     events = NanoEventsFactory.from_root(
-        path, schemaclass=PFNanoAODSchema, permit_dask=True
+        {path: "Events"}, schemaclass=PFNanoAODSchema, permit_dask=True
     ).events()
     return events
 

diff --git a/tests/test_nanoevents_physlite.py b/tests/test_nanoevents_physlite.py
@@ -11,8 +11,7 @@
 def _events():
     path = os.path.abspath("tests/samples/DAOD_PHYSLITE_21.2.108.0.art.pool.root")
     factory = NanoEventsFactory.from_root(
-        path,
-        treepath="CollectionTree",
+        {path: "CollectionTree"},
         schemaclass=PHYSLITESchema,
         permit_dask=False,
     )

diff --git a/tests/test_nanoevents_treemaker.py b/tests/test_nanoevents_treemaker.py
@@ -11,7 +11,7 @@
 def events():
     path = os.path.abspath("tests/samples/treemaker.root")
     events = NanoEventsFactory.from_root(
-        path, treepath="PreSelection", schemaclass=TreeMakerSchema, permit_dask=True
+        {path: "PreSelection"}, schemaclass=TreeMakerSchema, permit_dask=True
     ).events()
     return events