fsspec 2024.3.1__py3-none-any.whl → 2024.5.0__py3-none-any.whl

fsspec/implementations/tests/test_reference.py

@@ -0,0 +1,720 @@
+import json
+import os
+
+import pytest
+
+import fsspec
+from fsspec.implementations.local import LocalFileSystem
+from fsspec.implementations.reference import (
+    LazyReferenceMapper,
+    ReferenceFileSystem,
+    ReferenceNotReachable,
+)
+from fsspec.tests.conftest import data, realfile, reset_files, server, win  # noqa: F401
+
+
+def test_simple(server):  # noqa: F811
+    refs = {
+        "a": b"data",
+        "b": (realfile, 0, 5),
+        "c": (realfile, 1, 5),
+        "d": b"base64:aGVsbG8=",
+    }
+    h = fsspec.filesystem("http")
+    fs = fsspec.filesystem("reference", fo=refs, fs=h)
+
+    assert fs.cat("a") == b"data"
+    assert fs.cat("b") == data[:5]
+    assert fs.cat("c") == data[1 : 1 + 5]
+    assert fs.cat("d") == b"hello"
+    with fs.open("d", "rt") as f:
+        assert f.read(2) == "he"
+
+
+def test_target_options(m):
+    m.pipe("data/0", b"hello")
+    refs = {"a": ["memory://data/0"]}
+    fn = "memory://refs.json.gz"
+    with fsspec.open(fn, "wt", compression="gzip") as f:
+        json.dump(refs, f)
+
+    fs = fsspec.filesystem("reference", fo=fn, target_options={"compression": "gzip"})
+    assert fs.cat("a") == b"hello"
+
+
+def test_ls(server):  # noqa: F811
+    refs = {"a": b"data", "b": (realfile, 0, 5), "c/d": (realfile, 1, 6)}
+    h = fsspec.filesystem("http")
+    fs = fsspec.filesystem("reference", fo=refs, fs=h)
+
+    assert fs.ls("", detail=False) == ["a", "b", "c"]
+    assert {"name": "c", "type": "directory", "size": 0} in fs.ls("", detail=True)
+    assert fs.find("") == ["a", "b", "c/d"]
+    assert fs.find("", withdirs=True) == ["a", "b", "c", "c/d"]
+    assert fs.find("c", detail=True) == {
+        "c/d": {"name": "c/d", "size": 6, "type": "file"}
+    }
+
+
+def test_nested_dirs_ls():
+    # issue #1430
+    refs = {"a": "A", "B/C/b": "B", "B/C/d": "d", "B/_": "_"}
+    fs = fsspec.filesystem("reference", fo=refs)
+    assert len(fs.ls("")) == 2
+    assert {e["name"] for e in fs.ls("")} == {"a", "B"}
+    assert len(fs.ls("B")) == 2
+    assert {e["name"] for e in fs.ls("B")} == {"B/C", "B/_"}
+
+
+def test_info(server):  # noqa: F811
+    refs = {
+        "a": b"data",
+        "b": (realfile, 0, 5),
+        "c/d": (realfile, 1, 6),
+        "e": (realfile,),
+    }
+    h = fsspec.filesystem("http", headers={"give_length": "true", "head_ok": "true"})
+    fs = fsspec.filesystem("reference", fo=refs, fs=h)
+    assert fs.size("a") == 4
+    assert fs.size("b") == 5
+    assert fs.size("c/d") == 6
+    assert fs.info("e")["size"] == len(data)
+
+
+def test_mutable(server, m):
+    refs = {
+        "a": b"data",
+        "b": (realfile, 0, 5),
+        "c/d": (realfile, 1, 6),
+        "e": (realfile,),
+    }
+    h = fsspec.filesystem("http", headers={"give_length": "true", "head_ok": "true"})
+    fs = fsspec.filesystem("reference", fo=refs, fs=h)
+    fs.rm("a")
+    assert not fs.exists("a")
+
+    bin_data = b"bin data"
+    fs.pipe("aa", bin_data)
+    assert fs.cat("aa") == bin_data
+
+    fs.save_json("memory://refs.json")
+    assert m.exists("refs.json")
+
+    fs = fsspec.filesystem("reference", fo="memory://refs.json", remote_protocol="http")
+    assert not fs.exists("a")
+    assert fs.cat("aa") == bin_data
+
+
+def test_put_get(tmpdir):
+    d1 = f"{tmpdir}/d1"
+    os.mkdir(d1)
+    with open(f"{d1}/a", "wb") as f:
+        f.write(b"1")
+    with open(f"{d1}/b", "wb") as f:
+        f.write(b"2")
+    d2 = f"{tmpdir}/d2"
+
+    fs = fsspec.filesystem("reference", fo={}, remote_protocol="file")
+    fs.put(d1, "out", recursive=True)
+
+    fs.get("out", d2, recursive=True)
+    assert open(f"{d2}/a", "rb").read() == b"1"
+    assert open(f"{d2}/b", "rb").read() == b"2"
+
+
+def test_put_get_single(tmpdir):
+    d1 = f"{tmpdir}/f1"
+    d2 = f"{tmpdir}/f2"
+    with open(d1, "wb") as f:
+        f.write(b"1")
+
+    # skip instance cache since this is the same kwargs as previous test
+    fs = fsspec.filesystem(
+        "reference", fo={}, remote_protocol="file", skip_instance_cache=True
+    )
+    fs.put_file(d1, "out")
+
+    fs.get_file("out", d2)
+    assert open(d2, "rb").read() == b"1"
+    fs.pipe({"hi": b"data"})
+    assert fs.cat("hi") == b"data"
+
+
+def test_defaults(server):  # noqa: F811
+    refs = {"a": b"data", "b": (None, 0, 5)}
+    fs = fsspec.filesystem(
+        "reference",
+        fo=refs,
+        target_protocol="http",
+        target=realfile,
+        remote_protocol="http",
+    )
+
+    assert fs.cat("a") == b"data"
+    assert fs.cat("b") == data[:5]
+
+
+jdata = """{
+    "metadata": {
+        ".zattrs": {
+            "Conventions": "UGRID-0.9.0"
+        },
+        ".zgroup": {
+            "zarr_format": 2
+        },
+        "adcirc_mesh/.zarray": {
+            "chunks": [
+                1
+            ],
+            "dtype": "<i4",
+            "shape": [
+                1
+            ],
+            "zarr_format": 2
+        },
+        "adcirc_mesh/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "mesh"
+            ],
+            "cf_role": "mesh_topology"
+        },
+        "adcirc_mesh/.zchunkstore": {
+            "adcirc_mesh/0": {
+                "offset": 8928,
+                "size": 4
+            },
+            "source": {
+                "array_name": "/adcirc_mesh",
+                "uri": "https://url"
+            }
+        }
+    },
+    "zarr_consolidated_format": 1
+}
+"""
+
+
+def test_spec1_expand():
+    pytest.importorskip("jinja2")
+    in_data = {
+        "version": 1,
+        "templates": {"u": "server.domain/path", "f": "{{c}}"},
+        "gen": [
+            {
+                "key": "gen_key{{i}}",
+                "url": "http://{{u}}_{{i}}",
+                "offset": "{{(i + 1) * 1000}}",
+                "length": "1000",
+                "dimensions": {"i": {"stop": 5}},
+            },
+            {
+                "key": "gen_key{{i}}",
+                "url": "http://{{u}}_{{i}}",
+                "dimensions": {"i": {"start": 5, "stop": 7}},
+            },
+        ],
+        "refs": {
+            "key0": "data",
+            "key1": ["http://target_url", 10000, 100],
+            "key2": ["http://{{u}}", 10000, 100],
+            "key3": ["http://{{f(c='text')}}", 10000, 100],
+            "key4": ["http://target_url"],
+        },
+    }
+    fs = fsspec.filesystem(
+        "reference", fo=in_data, target_protocol="http", simple_templates=False
+    )
+    assert fs.references == {
+        "key0": "data",
+        "key1": ["http://target_url", 10000, 100],
+        "key2": ["http://server.domain/path", 10000, 100],
+        "key3": ["http://text", 10000, 100],
+        "key4": ["http://target_url"],
+        "gen_key0": ["http://server.domain/path_0", 1000, 1000],
+        "gen_key1": ["http://server.domain/path_1", 2000, 1000],
+        "gen_key2": ["http://server.domain/path_2", 3000, 1000],
+        "gen_key3": ["http://server.domain/path_3", 4000, 1000],
+        "gen_key4": ["http://server.domain/path_4", 5000, 1000],
+        "gen_key5": ["http://server.domain/path_5"],
+        "gen_key6": ["http://server.domain/path_6"],
+    }
+
+
+def test_spec1_expand_simple():
+    pytest.importorskip("jinja2")
+    in_data = {
+        "version": 1,
+        "templates": {"u": "server.domain/path"},
+        "refs": {
+            "key0": "base64:ZGF0YQ==",
+            "key2": ["http://{{u}}", 10000, 100],
+            "key4": ["http://target_url"],
+        },
+    }
+    fs = fsspec.filesystem("reference", fo=in_data, target_protocol="http")
+    assert fs.references["key2"] == ["http://server.domain/path", 10000, 100]
+    fs = fsspec.filesystem(
+        "reference",
+        fo=in_data,
+        target_protocol="http",
+        template_overrides={"u": "not.org/p"},
+    )
+    assert fs.references["key2"] == ["http://not.org/p", 10000, 100]
+    assert fs.cat("key0") == b"data"
+
+
+def test_spec1_gen_variants():
+    pytest.importorskip("jinja2")
+    with pytest.raises(ValueError):
+        missing_length_spec = {
+            "version": 1,
+            "templates": {"u": "server.domain/path"},
+            "gen": [
+                {
+                    "key": "gen_key{{i}}",
+                    "url": "http://{{u}}_{{i}}",
+                    "offset": "{{(i + 1) * 1000}}",
+                    "dimensions": {"i": {"stop": 2}},
+                },
+            ],
+        }
+        fsspec.filesystem("reference", fo=missing_length_spec, target_protocol="http")
+
+    with pytest.raises(ValueError):
+        missing_offset_spec = {
+            "version": 1,
+            "templates": {"u": "server.domain/path"},
+            "gen": [
+                {
+                    "key": "gen_key{{i}}",
+                    "url": "http://{{u}}_{{i}}",
+                    "length": "1000",
+                    "dimensions": {"i": {"stop": 2}},
+                },
+            ],
+        }
+        fsspec.filesystem("reference", fo=missing_offset_spec, target_protocol="http")
+
+    url_only_gen_spec = {
+        "version": 1,
+        "templates": {"u": "server.domain/path"},
+        "gen": [
+            {
+                "key": "gen_key{{i}}",
+                "url": "http://{{u}}_{{i}}",
+                "dimensions": {"i": {"stop": 2}},
+            },
+        ],
+    }
+
+    fs = fsspec.filesystem("reference", fo=url_only_gen_spec, target_protocol="http")
+    assert fs.references == {
+        "gen_key0": ["http://server.domain/path_0"],
+        "gen_key1": ["http://server.domain/path_1"],
+    }
+
+
+def test_empty():
+    pytest.importorskip("jinja2")
+    fs = fsspec.filesystem("reference", fo={"version": 1}, target_protocol="http")
+    assert fs.references == {}
+
+
+def test_get_sync(tmpdir):
+    localfs = LocalFileSystem()
+
+    real = tmpdir / "file"
+    real.write_binary(b"0123456789")
+
+    refs = {"a": b"data", "b": (str(real), 0, 5), "c/d": (str(real), 1, 6)}
+    fs = fsspec.filesystem("reference", fo=refs, fs=localfs)
+
+    fs.get("a", str(tmpdir / "a"))
+    assert (tmpdir / "a").read_binary() == b"data"
+    fs.get("b", str(tmpdir / "b"))
+    assert (tmpdir / "b").read_binary() == b"01234"
+    fs.get("c/d", str(tmpdir / "d"))
+    assert (tmpdir / "d").read_binary() == b"123456"
+    fs.get("c", str(tmpdir / "c"), recursive=True)
+    assert (tmpdir / "c").isdir()
+    assert (tmpdir / "c" / "d").read_binary() == b"123456"
+
+
+def test_multi_fs_provided(m, tmpdir):
+    localfs = LocalFileSystem()
+
+    real = tmpdir / "file"
+    real.write_binary(b"0123456789")
+
+    m.pipe("afile", b"hello")
+
+    # local URLs are file:// by default
+    refs = {
+        "a": b"data",
+        "b": (f"file://{real}", 0, 5),
+        "c/d": (f"file://{real}", 1, 6),
+        "c/e": ["memory://afile"],
+    }
+
+    fs = fsspec.filesystem("reference", fo=refs, fs={"file": localfs, "memory": m})
+    assert fs.cat("c/e") == b"hello"
+    assert fs.cat(["c/e", "a", "b"]) == {
+        "a": b"data",
+        "b": b"01234",
+        "c/e": b"hello",
+    }
+
+
+def test_multi_fs_created(m, tmpdir):
+    real = tmpdir / "file"
+    real.write_binary(b"0123456789")
+
+    m.pipe("afile", b"hello")
+
+    # local URLs are file:// by default
+    refs = {
+        "a": b"data",
+        "b": (f"file://{real}", 0, 5),
+        "c/d": (f"file://{real}", 1, 6),
+        "c/e": ["memory://afile"],
+    }
+
+    fs = fsspec.filesystem("reference", fo=refs, fs={"file": {}, "memory": {}})
+    assert fs.cat("c/e") == b"hello"
+    assert fs.cat(["c/e", "a", "b"]) == {
+        "a": b"data",
+        "b": b"01234",
+        "c/e": b"hello",
+    }
+
+
+def test_missing_nonasync(m):
+    zarr = pytest.importorskip("zarr")
+    zarray = {
+        "chunks": [1],
+        "compressor": None,
+        "dtype": "<f8",
+        "fill_value": "NaN",
+        "filters": [],
+        "order": "C",
+        "shape": [10],
+        "zarr_format": 2,
+    }
+    refs = {".zarray": json.dumps(zarray)}
+
+    m = fsspec.get_mapper("reference://", fo=refs, remote_protocol="memory")
+
+    a = zarr.open_array(m)
+    assert str(a[0]) == "nan"
+
+
+def test_fss_has_defaults(m):
+    fs = fsspec.filesystem("reference", fo={})
+    assert None in fs.fss
+
+    fs = fsspec.filesystem("reference", fo={}, remote_protocol="memory")
+    assert fs.fss[None].protocol == "memory"
+    assert fs.fss["memory"].protocol == "memory"
+
+    fs = fsspec.filesystem("reference", fs=m, fo={})
+    assert fs.fss[None] is m
+
+    fs = fsspec.filesystem("reference", fs={"memory": m}, fo={})
+    assert fs.fss["memory"] is m
+    assert fs.fss[None].protocol == ("file", "local")
+
+    fs = fsspec.filesystem("reference", fs={None: m}, fo={})
+    assert fs.fss[None] is m
+
+    fs = fsspec.filesystem("reference", fo={"key": ["memory://a"]})
+    assert fs.fss[None] is fs.fss["memory"]
+
+    fs = fsspec.filesystem("reference", fo={"key": ["memory://a"], "blah": ["path"]})
+    assert fs.fss[None] is fs.fss["memory"]
+
+
+def test_merging(m):
+    m.pipe("/a", b"test data")
+    other = b"other test data"
+    m.pipe("/b", other)
+    fs = fsspec.filesystem(
+        "reference",
+        fo={
+            "a": ["memory://a", 1, 1],
+            "b": ["memory://a", 2, 1],
+            "c": ["memory://b"],
+            "d": ["memory://b", 4, 6],
+        },
+    )
+    out = fs.cat(["a", "b", "c", "d"])
+    assert out == {"a": b"e", "b": b"s", "c": other, "d": other[4:10]}
+
+
+def test_cat_file_ranges(m):
+    other = b"other test data"
+    m.pipe("/b", other)
+
+    fs = fsspec.filesystem(
+        "reference",
+        fo={
+            "c": ["memory://b"],
+            "d": ["memory://b", 4, 6],
+        },
+    )
+    assert fs.cat_file("c") == other
+    assert fs.cat_file("c", start=1) == other[1:]
+    assert fs.cat_file("c", start=-5) == other[-5:]
+    assert fs.cat_file("c", 1, -5) == other[1:-5]
+
+    assert fs.cat_file("d") == other[4:10]
+    assert fs.cat_file("d", start=1) == other[4:10][1:]
+    assert fs.cat_file("d", start=-5) == other[4:10][-5:]
+    assert fs.cat_file("d", 1, -3) == other[4:10][1:-3]
+
+
+@pytest.mark.parametrize(
+    "fo",
+    [
+        {
+            "c": ["memory://b"],
+            "d": ["memory://unknown", 4, 6],
+        },
+        {
+            "c": ["memory://b"],
+            "d": ["//unknown", 4, 6],
+        },
+    ],
+    ids=["memory protocol", "mixed protocols: memory and unspecified"],
+)
+def test_cat_missing(m, fo):
+    other = b"other test data"
+    m.pipe("/b", other)
+    fs = fsspec.filesystem(
+        "reference",
+        fo=fo,
+    )
+    with pytest.raises(FileNotFoundError):
+        fs.cat("notafile")
+
+    with pytest.raises(FileNotFoundError):
+        fs.cat(["notone", "nottwo"])
+
+    mapper = fs.get_mapper("")
+
+    with pytest.raises(KeyError):
+        mapper["notakey"]
+
+    with pytest.raises(KeyError):
+        mapper.getitems(["notone", "nottwo"])
+
+    with pytest.raises(ReferenceNotReachable) as ex:
+        fs.cat("d")
+    assert ex.value.__cause__
+    out = fs.cat("d", on_error="return")
+    assert isinstance(out, ReferenceNotReachable)
+
+    with pytest.raises(ReferenceNotReachable) as e:
+        mapper["d"]
+    assert '"d"' in str(e.value)
+    assert "//unknown" in str(e.value)
+
+    with pytest.raises(ReferenceNotReachable):
+        mapper.getitems(["c", "d"])
+
+    out = mapper.getitems(["c", "d"], on_error="return")
+    assert isinstance(out["d"], ReferenceNotReachable)
+
+    out = fs.cat(["notone", "c", "d"], on_error="return")
+    assert isinstance(out["notone"], FileNotFoundError)
+    assert out["c"] == other
+    assert isinstance(out["d"], ReferenceNotReachable)
+
+    out = mapper.getitems(["c", "d"], on_error="omit")
+    assert list(out) == ["c"]
+
+
+def test_df_single(m):
+    pd = pytest.importorskip("pandas")
+    pytest.importorskip("fastparquet")
+    data = b"data0data1data2"
+    m.pipe({"data": data})
+    df = pd.DataFrame(
+        {
+            "path": [None, "memory://data", "memory://data"],
+            "offset": [0, 0, 4],
+            "size": [0, 0, 4],
+            "raw": [b"raw", None, None],
+        }
+    )
+    df.to_parquet("memory://stuff/refs.0.parq")
+    m.pipe(
+        ".zmetadata",
+        b"""{
+    "metadata": {
+        ".zgroup": {
+            "zarr_format": 2
+        },
+        "stuff/.zarray": {
+            "chunks": [1],
+            "compressor": null,
+            "dtype": "i8",
+            "filters": null,
+            "shape": [3],
+            "zarr_format": 2
+        }
+    },
+    "zarr_consolidated_format": 1,
+    "record_size": 10
+    }
+    """,
+    )
+    fs = ReferenceFileSystem(fo="memory:///", remote_protocol="memory")
+    allfiles = fs.find("")
+    assert ".zmetadata" in allfiles
+    assert ".zgroup" in allfiles
+    assert "stuff/2" in allfiles
+
+    assert fs.cat("stuff/0") == b"raw"
+    assert fs.cat("stuff/1") == data
+    assert fs.cat("stuff/2") == data[4:8]
+
+
+def test_df_multi(m):
+    pd = pytest.importorskip("pandas")
+    pytest.importorskip("fastparquet")
+    data = b"data0data1data2"
+    m.pipe({"data": data})
+    df0 = pd.DataFrame(
+        {
+            "path": [None, "memory://data", "memory://data"],
+            "offset": [0, 0, 4],
+            "size": [0, 0, 4],
+            "raw": [b"raw1", None, None],
+        }
+    )
+    df0.to_parquet("memory://stuff/refs.0.parq")
+    df1 = pd.DataFrame(
+        {
+            "path": [None, "memory://data", "memory://data"],
+            "offset": [0, 0, 2],
+            "size": [0, 0, 2],
+            "raw": [b"raw2", None, None],
+        }
+    )
+    df1.to_parquet("memory://stuff/refs.1.parq")
+    m.pipe(
+        ".zmetadata",
+        b"""{
+    "metadata": {
+        ".zgroup": {
+            "zarr_format": 2
+        },
+        "stuff/.zarray": {
+            "chunks": [1],
+            "compressor": null,
+            "dtype": "i8",
+            "filters": null,
+            "shape": [6],
+            "zarr_format": 2
+        }
+    },
+    "zarr_consolidated_format": 1,
+    "record_size": 3
+    }
+    """,
+    )
+    fs = ReferenceFileSystem(
+        fo="memory:///", remote_protocol="memory", skip_instance_cache=True
+    )
+    allfiles = fs.find("")
+    assert ".zmetadata" in allfiles
+    assert ".zgroup" in allfiles
+    assert "stuff/2" in allfiles
+    assert "stuff/4" in allfiles
+
+    assert fs.cat("stuff/0") == b"raw1"
+    assert fs.cat("stuff/1") == data
+    assert fs.cat("stuff/2") == data[4:8]
+    assert fs.cat("stuff/3") == b"raw2"
+    assert fs.cat("stuff/4") == data
+    assert fs.cat("stuff/5") == data[2:4]
+
+
+def test_mapping_getitems(m):
+    m.pipe({"a": b"A", "b": b"B"})
+
+    refs = {
+        "a": ["a"],
+        "b": ["b"],
+    }
+    h = fsspec.filesystem("memory")
+    fs = fsspec.filesystem("reference", fo=refs, fs=h)
+    mapping = fs.get_mapper("")
+    assert mapping.getitems(["b", "a"]) == {"a": b"A", "b": b"B"}
+
+
+def test_cached(m, tmpdir):
+    fn = f"{tmpdir}/ref.json"
+
+    m.pipe({"a": b"A", "b": b"B"})
+    m.pipe("ref.json", b"""{"a": ["a"], "b": ["b"]}""")
+
+    fs = fsspec.filesystem(
+        "reference",
+        fo="simplecache::memory://ref.json",
+        fs=m,
+        target_options={"cache_storage": str(tmpdir), "same_names": True},
+    )
+    assert fs.cat("a") == b"A"
+    assert os.path.exists(fn)
+
+    # truncate original file to show we are loading from the cached version
+    m.pipe("ref.json", b"")
+    fs = fsspec.filesystem(
+        "reference",
+        fo="simplecache::memory://ref.json",
+        fs=m,
+        target_options={"cache_storage": str(tmpdir), "same_names": True},
+        skip_instance_cache=True,
+    )
+    assert fs.cat("a") == b"A"
+
+
+@pytest.fixture()
+def lazy_refs(m):
+    zarr = pytest.importorskip("zarr")
+    l = LazyReferenceMapper.create("memory://refs", fs=m)
+    g = zarr.open(l, mode="w")
+    g.create_dataset(name="data", shape=(100,), chunks=(10,), dtype="int64")
+    return l
+
+
+def test_append_parquet(lazy_refs, m):
+    pytest.importorskip("kerchunk")
+    with pytest.raises(KeyError):
+        lazy_refs["data/0"]
+    lazy_refs["data/0"] = b"data"
+    assert lazy_refs["data/0"] == b"data"
+    lazy_refs.flush()
+
+    lazy2 = LazyReferenceMapper("memory://refs", fs=m)
+    assert lazy2["data/0"] == b"data"
+    with pytest.raises(KeyError):
+        lazy_refs["data/1"]
+    lazy2["data/1"] = b"Bdata"
+    assert lazy2["data/1"] == b"Bdata"
+    lazy2.flush()
+
+    lazy2 = LazyReferenceMapper("memory://refs", fs=m)
+    assert lazy2["data/0"] == b"data"
+    assert lazy2["data/1"] == b"Bdata"
+    lazy2["data/1"] = b"Adata"
+    del lazy2["data/0"]
+    assert lazy2["data/1"] == b"Adata"
+    assert "data/0" not in lazy2
+    lazy2.flush()
+
+    lazy2 = LazyReferenceMapper("memory://refs", fs=m)
+    with pytest.raises(KeyError):
+        lazy2["data/0"]
+    assert lazy2["data/1"] == b"Adata"