Skip to content

Commit

Permalink
fixed kerchunk tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@TomNicholas
TomNicholas committed Mar 18, 2024
1 parent bd8ad22 commit 385290d
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 14 deletions.
6 changes: 3 additions & 3 deletions virtualizarr/kerchunk.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ def variable_to_kerchunk_arr_refs(var: xr.Variable) -> KerchunkArrRefs:
Partially encodes the inner dicts to json to match kerchunk behaviour (see https://github.com/fsspec/kerchunk/issues/415).
"""
from virtualizarr.manifests import ManifestArray
from virtualizarr.manifests import ChunkEntry, ManifestArray

marr = var.data

Expand All @@ -177,8 +177,8 @@ def variable_to_kerchunk_arr_refs(var: xr.Variable) -> KerchunkArrRefs:
)

arr_refs: dict[str, Union[str, List[Union[str, int]]]] = {
str(chunk_key): chunk_entry.to_kerchunk()
for chunk_key, chunk_entry in marr.manifest.entries.items()
str(chunk_key): ChunkEntry(**chunk_entry).to_kerchunk()
for chunk_key, chunk_entry in marr.manifest.dict().items()
}

zarray_dict = marr.zarray.to_kerchunk_json()
Expand Down
14 changes: 10 additions & 4 deletions virtualizarr/tests/test_kerchunk.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import xarray as xr
import xarray.testing as xrt

from virtualizarr.manifests import ChunkEntry, ChunkManifest, ManifestArray
from virtualizarr.manifests import ChunkManifest, ManifestArray
from virtualizarr.xarray import dataset_from_kerchunk_refs


Expand Down Expand Up @@ -40,8 +40,11 @@ def test_dataset_from_kerchunk_refs():

class TestAccessor:
def test_accessor_to_kerchunk_dict(self):
manifest = ChunkManifest.from_dict(
{"0.0": dict(path="test.nc", offset=6144, length=48)}
)
arr = ManifestArray(
chunkmanifest={"0.0": ChunkEntry(path="test.nc", offset=6144, length=48)},
chunkmanifest=manifest,
zarray=dict(
shape=(2, 3),
dtype=np.dtype("<i8"),
Expand All @@ -68,8 +71,11 @@ def test_accessor_to_kerchunk_dict(self):
assert result_ds_refs == expected_ds_refs

def test_accessor_to_kerchunk_json(self, tmp_path):
manifest = ChunkManifest.from_dict(
{"0.0": dict(path="test.nc", offset=6144, length=48)}
)
arr = ManifestArray(
chunkmanifest={"0.0": ChunkEntry(path="test.nc", offset=6144, length=48)},
chunkmanifest=manifest,
zarray=dict(
shape=(2, 3),
dtype=np.dtype("<i8"),
Expand Down Expand Up @@ -107,7 +113,7 @@ def test_kerchunk_roundtrip_in_memory_no_concat():
"0.0": {"path": "foo.nc", "offset": 100, "length": 100},
"0.1": {"path": "foo.nc", "offset": 200, "length": 100},
}
manifest = ChunkManifest(entries=chunks_dict)
manifest = ChunkManifest.from_dict(chunks_dict)
marr = ManifestArray(
zarray=dict(
shape=(2, 4),
Expand Down
14 changes: 7 additions & 7 deletions virtualizarr/tests/test_xarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def test_wrapping():
"0.0": {"path": "foo.nc", "offset": 100, "length": 100},
"0.1": {"path": "foo.nc", "offset": 200, "length": 100},
}
manifest = ChunkManifest(entries=chunks_dict)
manifest = ChunkManifest.from_dict(chunks_dict)
marr = ManifestArray(zarray=zarray, chunkmanifest=manifest)
ds = xr.Dataset({"a": (["x", "y"], marr)})

Expand Down Expand Up @@ -54,7 +54,7 @@ def test_equals(self):
"0.0": {"path": "foo.nc", "offset": 100, "length": 100},
"0.1": {"path": "foo.nc", "offset": 200, "length": 100},
}
manifest1 = ChunkManifest(entries=chunks_dict1)
manifest1 = ChunkManifest.from_dict(chunks_dict1)
marr1 = ManifestArray(zarray=zarray, chunkmanifest=manifest1)
ds1 = xr.Dataset({"a": (["x", "y"], marr1)})

Expand All @@ -66,7 +66,7 @@ def test_equals(self):
"0.0": {"path": "foo.nc", "offset": 300, "length": 100},
"0.1": {"path": "foo.nc", "offset": 400, "length": 100},
}
manifest3 = ChunkManifest(entries=chunks_dict3)
manifest3 = ChunkManifest.from_dict(chunks_dict3)
marr3 = ManifestArray(zarray=zarray, chunkmanifest=manifest3)
ds3 = xr.Dataset({"a": (["x", "y"], marr3)})
assert not ds1.equals(ds3)
Expand All @@ -90,15 +90,15 @@ def test_concat_along_existing_dim(self):
"0.0": {"path": "foo.nc", "offset": 100, "length": 100},
"0.1": {"path": "foo.nc", "offset": 200, "length": 100},
}
manifest1 = ChunkManifest(entries=chunks_dict1)
manifest1 = ChunkManifest.from_dict(chunks_dict1)
marr1 = ManifestArray(zarray=zarray, chunkmanifest=manifest1)
ds1 = xr.Dataset({"a": (["x", "y"], marr1)})

chunks_dict2 = {
"0.0": {"path": "foo.nc", "offset": 300, "length": 100},
"0.1": {"path": "foo.nc", "offset": 400, "length": 100},
}
manifest2 = ChunkManifest(entries=chunks_dict2)
manifest2 = ChunkManifest.from_dict(chunks_dict2)
marr2 = ManifestArray(zarray=zarray, chunkmanifest=manifest2)
ds2 = xr.Dataset({"a": (["x", "y"], marr2)})

Expand Down Expand Up @@ -136,15 +136,15 @@ def test_concat_along_new_dim(self):
"0.0": {"path": "foo.nc", "offset": 100, "length": 100},
"0.1": {"path": "foo.nc", "offset": 200, "length": 100},
}
manifest1 = ChunkManifest(entries=chunks_dict1)
manifest1 = ChunkManifest.from_dict(chunks_dict1)
marr1 = ManifestArray(zarray=zarray, chunkmanifest=manifest1)
ds1 = xr.Dataset({"a": (["x", "y"], marr1)})

chunks_dict2 = {
"0.0": {"path": "foo.nc", "offset": 300, "length": 100},
"0.1": {"path": "foo.nc", "offset": 400, "length": 100},
}
manifest2 = ChunkManifest(entries=chunks_dict2)
manifest2 = ChunkManifest.from_dict(chunks_dict2)
marr2 = ManifestArray(zarray=zarray, chunkmanifest=manifest2)
ds2 = xr.Dataset({"a": (["x", "y"], marr2)})

Expand Down

0 comments on commit 385290d

Please sign in to comment.