From 2b402c02521d19045031030efb0e5a4b7898eaa6 Mon Sep 17 00:00:00 2001
From: Dan Allan <dallan@bnl.gov>
Date: Mon, 12 Feb 2024 09:29:05 -0500
Subject: [PATCH] Xarray specs bug (#653)

* Test xarray specs mutation bug.

* Do not mutate user input.

* Append 'xarray_dataset' if not included in specs.
---
 tiled/_tests/test_xarray.py | 45 +++++++++++++++++++++++++++++++++++++
 tiled/adapters/xarray.py    |  5 +++--
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/tiled/_tests/test_xarray.py b/tiled/_tests/test_xarray.py
index 5a9299e9b..27e8416e7 100644
--- a/tiled/_tests/test_xarray.py
+++ b/tiled/_tests/test_xarray.py
@@ -1,6 +1,7 @@
 import dask.array
 import numpy
 import orjson
+import pandas
 import pytest
 import xarray
 import xarray.testing
@@ -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)
@@ -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]
diff --git a/tiled/adapters/xarray.py b/tiled/adapters/xarray.py
index 710248dae..f8e9088c1 100644
--- a/tiled/adapters/xarray.py
+++ b/tiled/adapters/xarray.py
@@ -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},
@@ -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
         )