Skip to content

Commit

Permalink
Xarray specs bug (#653)
Browse files Browse the repository at this point in the history
* Test xarray specs mutation bug.

* Do not mutate user input.

* Append 'xarray_dataset' if not included in specs.
  • Loading branch information
danielballan authored Feb 12, 2024
1 parent 23b18f2 commit 2b402c0
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 2 deletions.
45 changes: 45 additions & 0 deletions tiled/_tests/test_xarray.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import dask.array
import numpy
import orjson
import pandas
import pytest
import xarray
import xarray.testing
Expand All @@ -11,6 +12,7 @@
from ..client import xarray as xarray_client
from ..serialization.xarray import serialize_json
from ..server.app import build_app
from ..structures.core import Spec

image = numpy.random.random((3, 5))
temp = 15 + 8 * numpy.random.randn(2, 2, 3)
Expand Down Expand Up @@ -79,6 +81,49 @@ def test_xarray_dataset(client, key):
xarray.testing.assert_equal(actual, expected)


def test_specs(client):
assert client["image"].specs == [Spec("xarray_dataset")]
assert client["image"]["image"].specs == [Spec("xarray_data_var")]
assert client["image"]["x"].specs == [Spec("xarray_coord")]


def test_specs_mutation_bug(client):
# https://github.com/bluesky/tiled/issues/651
ds = pandas.DataFrame({"x": numpy.array([1, 2, 3])}).to_xarray()
tree = MapAdapter({"data": DatasetAdapter.from_dataset(ds)})
for _ in range(2):
# This bug caused additional, redundant specs to be appended on each
# iteration.
app = build_app(tree)
with Context.from_app(app) as context:
client = from_context(context)
data = client["data"]
data.read()
assert data.specs == [Spec("xarray_dataset")]


def test_specs_override(client):
"The 'xarray_dataset' is appended to the end if not present."
ds = pandas.DataFrame({"x": numpy.array([1, 2, 3])}).to_xarray()
tree = MapAdapter(
{
"a": DatasetAdapter.from_dataset(ds, specs=[Spec("test")]),
"b": DatasetAdapter.from_dataset(
ds, specs=[Spec("xarray_dataset"), Spec("test")]
),
"c": DatasetAdapter.from_dataset(
ds, specs=[Spec("test"), Spec("xarray_dataset")]
),
}
)
app = build_app(tree)
with Context.from_app(app) as context:
client = from_context(context)
assert client["a"].specs == [Spec("test"), Spec("xarray_dataset")]
assert client["b"].specs == [Spec("xarray_dataset"), Spec("test")]
assert client["c"].specs == [Spec("test"), Spec("xarray_dataset")]


@pytest.mark.parametrize("key", ["image", "weather"])
def test_dataset_column_access(client, key):
expected_dataset = EXPECTED[key]
Expand Down
5 changes: 3 additions & 2 deletions tiled/adapters/xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@ class DatasetAdapter(MapAdapter):
@classmethod
def from_dataset(cls, dataset, *, specs=None, access_policy=None):
mapping = _DatasetMap(dataset)
specs = specs or []
if "xarray_dataset" not in [spec.name for spec in specs]:
specs.append(Spec("xarray_dataset"))
return cls(
mapping,
metadata={"attrs": dataset.attrs},
Expand All @@ -28,8 +31,6 @@ def __init__(self, mapping, *args, specs=None, access_policy=None, **kwargs):
raise TypeError(
"Use DatasetAdapter.from_dataset(...), not DatasetAdapter(...)."
)
specs = specs or []
specs.append(Spec("xarray_dataset"))
super().__init__(
mapping, *args, specs=specs, access_policy=access_policy, **kwargs
)
Expand Down

0 comments on commit 2b402c0

Please sign in to comment.