Skip to content

Commit

Permalink
feat: correct fsspec source serialization (#1033)
Browse files Browse the repository at this point in the history
* test different handlers

* correct serialization of fsspec source

* fsspec is not yet required
  • Loading branch information
lobis authored Nov 16, 2023
1 parent 0b5d8d6 commit c956761
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 13 deletions.
15 changes: 12 additions & 3 deletions src/uproot/source/fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,20 +33,27 @@ def __init__(self, file_path: str, **options):
if k not in uproot.reading.open.defaults.keys()
}

self._executor = FSSpecLoopExecutor()

self._fs, self._file_path = fsspec.core.url_to_fs(file_path, **storage_options)

# What should we do when there is a chain of filesystems?
self._async_impl = self._fs.async_impl

self._file = self._fs.open(self._file_path)
self._executor = None
self._file = None
self._fh = None

self._num_requests = 0
self._num_requested_chunks = 0
self._num_requested_bytes = 0

self._open()

self.__enter__()

def _open(self):
self._executor = FSSpecLoopExecutor()
self._file = self._fs.open(self._file_path)

def __repr__(self):
path = repr(self._file_path)
if len(self._file_path) > 10:
Expand All @@ -56,6 +63,8 @@ def __repr__(self):
def __getstate__(self):
state = dict(self.__dict__)
state.pop("_executor")
state.pop("_file")
state.pop("_fh")
return state

def __setstate__(self, state):
Expand Down
44 changes: 34 additions & 10 deletions tests/test_0302_pickle.py
Original file line number Diff line number Diff line change
@@ -1,19 +1,24 @@
# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE

import os
import pickle
import sys

import pytest
import skhep_testdata

import uproot

pytest.importorskip("awkward")


def test_pickle_roundtrip_mmap():
with uproot.open(skhep_testdata.data_path("uproot-small-dy-withoffsets.root")) as f:
@pytest.mark.parametrize(
"handler",
[
uproot.source.file.MemmapSource,
# uproot.source.fsspec.FSSpecSource,
],
)
def test_pickle_roundtrip_local(handler):
with uproot.open(
skhep_testdata.data_path("uproot-small-dy-withoffsets.root"), handler=handler
) as f:
pkl = pickle.dumps(f["tree"])

branch = pickle.loads(pkl)["Muon_pt"]
Expand All @@ -32,9 +37,20 @@ def test_pickle_roundtrip_mmap():
]


@pytest.mark.parametrize(
"handler",
[
uproot.source.http.HTTPSource,
# uproot.source.fsspec.FSSpecSource,
],
)
@pytest.mark.network
def test_pickle_roundtrip_http():
with uproot.open("https://scikit-hep.org/uproot3/examples/Zmumu.root") as f:
def test_pickle_roundtrip_http(handler):
pytest.importorskip("aiohttp")

with uproot.open(
"https://scikit-hep.org/uproot3/examples/Zmumu.root", handler=handler
) as f:
pkl = pickle.dumps(f["events"])

tree = pickle.loads(pkl)
Expand All @@ -53,12 +69,20 @@ def test_pickle_roundtrip_http():
]


@pytest.mark.parametrize(
"handler",
[
uproot.source.xrootd.XRootDSource,
# uproot.source.fsspec.FSSpecSource,
],
)
@pytest.mark.network
@pytest.mark.xrootd
def test_pickle_roundtrip_xrootd():
def test_pickle_roundtrip_xrootd(handler):
pytest.importorskip("XRootD")
with uproot.open(
"root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012B_DoubleMuParked.root"
"root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012B_DoubleMuParked.root",
handler=handler,
) as f:
pkl = pickle.dumps(f["Events"])

Expand Down

0 comments on commit c956761

Please sign in to comment.