🩹 Fix find_axes inside use_plot_config causing TypeError with DataArr…

…ay (#303)

* 🧪 Added test reproducing error in find_axes when a xr.DataArray is passed
* 🩹 Fix issue with xr.DataArray
* 👌 Endure Mapping of Axes is handled correctly by find_axes
* 🚧📚 Added change to changelog

changelog.md

@@ -4,6 +4,8 @@
 
 ## 0.7.4 (Unreleased)
 
+- 🩹 Fix find_axes inside use_plot_config causing TypeError with DataArray (#303)
+
 (changes-0_7_3)=
 
 ## 0.7.3 (2024-08-25)

pyglotaran_extras/config/plot_config.py

@@ -20,6 +20,7 @@
 from typing import cast
 
 import numpy as np
+import xarray as xr
 from docstring_parser import parse as parse_docstring
 from matplotlib.axes import Axes
 from pydantic import BaseModel
@@ -479,12 +480,16 @@ def find_axes(
     Axes
     """
     for value in values:
-        if isinstance(value, str):
+        # This are iterables where we are sure that they can not contain `Axes` so we can skip them
+        # early
+        if isinstance(value, str | xr.Dataset | xr.DataArray):
             continue
         elif isinstance(value, Axes):
             yield value
         elif isinstance(value, np.ndarray):
             yield from find_axes(value.flatten())
+        elif isinstance(value, Mapping):
+            yield from find_axes(value.values())
         elif isinstance(value, Iterable):
             yield from find_axes(value)
 

tests/config/test_plot_config.py

@@ -10,6 +10,7 @@
 
 import matplotlib.pyplot as plt
 import pytest
+import xarray as xr
 from jsonschema import ValidationError as SchemaValidationError
 from jsonschema import validate
 from pydantic import ValidationError as PydanticValidationError
@@ -466,15 +467,18 @@ def func(
 def test_find_axes():
     """Get axes value from iterable of values."""
 
-    base_values = ["foo", True, 1.5]
+    data_array = xr.DataArray([[0, 1]], coords={"time": [0], "spectral": [500, 510]})
+    data_set = xr.Dataset({"data": data_array})
+
+    base_values = ["foo", True, 1.5, data_array, data_set]
 
     assert generator_is_exhausted(find_axes(base_values)) is True
 
     _, ax = plt.subplots()
-    single_ax_gen = find_axes([*base_values, ax])
+    dict_ax_gen = find_axes([*base_values, ax])
 
-    assert next(single_ax_gen) is ax
-    assert generator_is_exhausted(single_ax_gen) is True
+    assert next(dict_ax_gen) is ax
+    assert generator_is_exhausted(dict_ax_gen) is True
 
     _, np_axes = plt.subplots(1, 2)
 
@@ -500,6 +504,11 @@ def test_find_axes():
     assert next(multiple_axes_gen) is ax1
     assert generator_is_exhausted(multiple_axes_gen) is True
 
+    dict_ax_gen = find_axes([*base_values, {"ax": ax}])
+
+    assert next(dict_ax_gen) is ax
+    assert generator_is_exhausted(dict_ax_gen) is True
+
 
 def test_use_plot_config(mock_config: tuple[Config, dict[str, Any]]):
     """Config is applied to functions with the ``use_plot_config`` decorator."""