feat: add file handle as an attribute of the events factory (#1502)

* feat: add file.source as metadata in processor

* add flag to only run filesource test in executor workflows

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* track file handle instead of fsspecsource and delete reference before exiting file context

* remove pixi dust added by mistake

* add file handle to events factory instead of metadata, following implementaiton for access_log

* fix precommit for test file

* remove extra del (nicks comment)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Iason Krommydas <iason.krom@gmail.com>

Author: 3 people

src/coffea/nanoevents/factory.py

@@ -392,13 +392,15 @@ def from_root(
 
         if isinstance(file, uproot.reading.ReadOnlyDirectory):
             tree = file[treepath]
+            file_handle = file
         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 uproot5 or higher"
                 % file
             )
         else:
             tree = uproot.open(file, **uproot_options)
+            file_handle = tree.file
 
         # Get the typenames
         typenames = tree.typenames()
@@ -438,6 +440,7 @@ def from_root(
             entry_stop,
             cache={},
             access_log=access_log,
+            file_handle=file_handle,
             use_ak_forth=use_ak_forth,
             virtual=mode == "virtual",
             preloaded_arrays=preloaded_arrays,
@@ -755,6 +758,11 @@ def access_log(self):
         """List of accessed branches, populated when columns are lazily loaded."""
         return getattr(self._mapping, "_access_log", None)
 
+    @property
+    def file_handle(self):
+        """The file handle used to open the source file, if available."""
+        return getattr(self._mapping, "_file_handle", None)
+
     def events(self):
         """
         Build events

src/coffea/nanoevents/mapping/base.py

@@ -34,13 +34,15 @@ def __init__(
         stop,
         cache=None,
         access_log=None,
+        file_handle=None,
         use_ak_forth=False,
         virtual=False,
         buffer_cache=None,
     ):
         self._fileopener = fileopener
         self._cache = cache
         self._access_log = access_log
+        self._file_handle = file_handle
         self._start = start
         self._stop = stop
         self._use_ak_forth = use_ak_forth

src/coffea/nanoevents/mapping/uproot.py

@@ -123,6 +123,7 @@ def __init__(
         stop,
         cache=None,
         access_log=None,
+        file_handle=None,
         use_ak_forth=False,
         virtual=False,
         decompression_executor=None,
@@ -136,6 +137,7 @@ def __init__(
             stop=stop,
             cache=cache,
             access_log=access_log,
+            file_handle=file_handle,
             use_ak_forth=use_ak_forth,
             virtual=virtual,
             buffer_cache=buffer_cache,

src/coffea/processor/executor.py

@@ -1550,6 +1550,7 @@ def _work_function(
                     )
                 # save the output
                 checkpointer.save(out, metadata, processor_instance)
+
             return out
 
     def __call__(

src/coffea/processor/test_items/NanoEventsProcessor.py

@@ -6,10 +6,11 @@
 
 
 class NanoEventsProcessor(processor.ProcessorABC):
-    def __init__(self, columns=[], canaries=[], mode="dask"):
+    def __init__(self, columns=[], canaries=[], mode="dask", check_filehandle=False):
         self._columns = columns
         self._canaries = canaries
         self.mode = mode
+        self.check_filehandle = check_filehandle
 
         self.expected_usermeta = {
             "ZJets": ("someusermeta", "hello"),
@@ -56,6 +57,10 @@ def process(self, events):
             metaname, metavalue = self.expected_usermeta[dataset]
             assert metavalue == events.metadata[metaname]
 
+        if self.check_filehandle:
+            factory = events.attrs["@events_factory"]
+            assert factory.file_handle is not None
+
         # mapping = events.behavior["__events_factory__"]._mapping
         muon_pt = events.Muon.pt
         # if isinstance(mapping, nanoevents.mapping.CachedMapping):

tests/test_local_executors.py

@@ -23,9 +23,11 @@ def test_nanoevents_analysis(
     executor, compression, maxchunks, skipbadfiles, filetype, mode, processor_type
 ):
     if processor_type == "Callable":
-        processor_instance = NanoEventsProcessor(mode=mode)
+        processor_instance = NanoEventsProcessor(mode=mode, check_filehandle=True)
     else:
-        processor_instance = NanoEventsProcessor(mode=mode).process
+        processor_instance = NanoEventsProcessor(
+            mode=mode, check_filehandle=True
+        ).process
 
     if filetype == "parquet":
         pytest.xfail("parquet nanoevents not supported yet")

tests/test_nanoevents.py

@@ -2,6 +2,7 @@
 
 import awkward as ak
 import pytest
+import uproot
 from distributed import Client
 
 from coffea.nanoevents import NanoAODSchema, NanoEventsFactory
@@ -246,3 +247,47 @@ def test_access_log(tests_directory, mode):
         _ = ak.materialize(events.Muon.pt)
         branches = {entry.branch for entry in factory.access_log}
         assert branches == {"nMuon", "Muon_pt"}
+
+
+@pytest.mark.parametrize("mode", ["eager", "virtual"])
+def test_file_handle_from_path(tests_directory, mode):
+    """Test that file_handle is available when opening from path string."""
+    path = f"{tests_directory}/samples/nano_dy.root:Events"
+
+    factory = NanoEventsFactory.from_root(
+        path,
+        schemaclass=NanoAODSchema,
+        mode=mode,
+    )
+
+    # file_handle should be ReadOnlyFile when opened from path
+    assert factory.file_handle is not None
+    assert isinstance(factory.file_handle, uproot.reading.ReadOnlyFile)
+
+    _ = factory.events()
+
+    # file_handle still accessible after events() call
+    assert factory.file_handle is not None
+
+
+@pytest.mark.parametrize("mode", ["eager", "virtual"])
+def test_file_handle_from_directory(tests_directory, mode):
+    """Test that file_handle is available when passing ReadOnlyDirectory."""
+    filepath = f"{tests_directory}/samples/nano_dy.root"
+
+    with uproot.open(filepath) as file:
+        factory = NanoEventsFactory.from_root(
+            file,
+            treepath="Events",
+            schemaclass=NanoAODSchema,
+            mode=mode,
+        )
+
+        # file_handle should be ReadOnlyDirectory when passed directly
+        assert factory.file_handle is not None
+        assert isinstance(factory.file_handle, uproot.ReadOnlyDirectory)
+
+        _ = factory.events()
+
+        # file_handle still accessible after events() call
+        assert factory.file_handle is not None

tests/test_skyhook_job.py

@@ -69,7 +69,7 @@
             "Data": "/mnt/cephfs/nanoevents/Data",
         },
         "Events",
-        processor_instance=NanoEventsProcessor(),
+        processor_instance=NanoEventsProcessor(check_filehandle=True),
     )
 
     assert hists["cutflow"]["ZJets_pt"] == 108
@@ -93,7 +93,7 @@
             "Data": "/mnt/cephfs/nanoevents/Data",
         },
         "Events",
-        processor_instance=NanoEventsProcessor(),
+        processor_instance=NanoEventsProcessor(check_filehandle=True),
     )
 
     assert hists["cutflow"]["ZJets_pt"] == 108

tests/test_taskvine_virtual.py

@@ -31,7 +31,7 @@ def test_taskvine_executor_with_virtual_arrays():
     }
 
     # Create the same processor as used in local executors
-    processor = NanoEventsProcessor(mode="virtual")
+    processor = NanoEventsProcessor(mode="virtual", check_filehandle=True)
 
     # Create Runner with TaskVineExecutor
     from coffea.processor import Runner