fix attribute serialization

Author: forman

tests/config/test_attrs.py

@@ -7,6 +7,7 @@
 from datetime import datetime
 
 import numpy as np
+import xarray as xr
 import pytest
 
 from zappend.config.attrs import ConfigAttrsUserFunctions
@@ -149,51 +150,56 @@ def test_yes(self):
 
 class EvalExprTest(unittest.TestCase):
     def test_all_cases(self):
+        # scalars
+        self.assertEqual(None, eval_expr("None", {}))
         self.assertEqual(True, eval_expr("True", {}))
         self.assertEqual(13, eval_expr("13", {}))
         self.assertEqual(0.5, eval_expr("0.5", {}))
         self.assertEqual("ABC", eval_expr("'ABC'", {}))
-
-        switches = np.array([True, False])
+        time = datetime.fromisoformat("2024-01-02T10:20:30")
         self.assertEqual(
-            True,
-            eval_expr("switches[0]", {"switches": switches}),
+            "2024-01-02T10:20:30",
+            eval_expr("time", dict(time=time)),
         )
 
-        levels = [3, 4, 5]
-        self.assertIs(
-            levels,
-            eval_expr("levels", {"levels": levels}),
-        )
+        # arrays
+        self.assert_array_ok([True, False])
+        self.assert_array_ok([3, 4, 5])
+        self.assert_array_ok([11.05, 11.15, 11.25])
+        self.assert_array_ok(["A", "B"])
+        self.assert_array_ok(["2024-01-02T10:20:30"], dtype="datetime64[s]")
 
-        levels = np.array([3, 4, 5])
+    def assert_array_ok(self, a: list, dtype=None):
+        # Test list
         self.assertEqual(
-            3,
-            eval_expr("levels[0]", {"levels": levels}),
+            a,
+            eval_expr("a", dict(a=a)),
         )
-
-        lon = np.array([11.05, 11.15, 11.25])
-        self.assertIs(
-            lon,
-            eval_expr("lon", {"lon": lon}),
+        self.assertEqual(
+            a[0],
+            eval_expr("a[0]", dict(a=a)),
         )
 
-        lon = np.array([11.05, 11.15, 11.25])
+        # Test numpy.ndarray
+        np_a = np.array(a, dtype=dtype) if dtype is not None else np.array(a)
         self.assertEqual(
-            11.05,
-            eval_expr("lon[0]", {"lon": lon}),
+            a,
+            eval_expr("a", dict(a=np_a)),
         )
-
-        names = np.array(["A", "B"])
         self.assertEqual(
-            "A",
-            eval_expr("names[0]", {"names": names}),
+            a[0],
+            eval_expr("a[0]", dict(a=np_a)),
         )
 
-        time = datetime.fromisoformat("2024-01-02T10:20:30")
+        # Test xarray-DataArray
+        xr_a = xr.DataArray(np_a, dims="x")
         self.assertEqual(
-            "2024-01-02T10:20:30",
-            eval_expr("time", {"time": time}),
+            a,
+            eval_expr("a", dict(a=xr_a)),
+        )
+        self.assertEqual(
+            a[0],
+            eval_expr("a[0]", dict(a=xr_a)),
         )
 
 

zappend/config/attrs.py

@@ -3,6 +3,7 @@
 # https://opensource.org/licenses/MIT.
 
 import datetime
+import math
 from typing import Any, Literal
 
 import numpy as np
@@ -64,24 +65,58 @@ def eval_attr_value(attr_value: Any, env: dict[str, Any]) -> Any:
 
 def eval_expr(expr: str, env: dict[str, Any]) -> Any:
     value = eval(expr, env)
-    if isinstance(value, (bool, int, float, str, type(None))):
+    return to_json(value)
+
+
+def to_json(value) -> Any:
+    if isinstance(value, float):
+        if math.isfinite(value):
+            return value
+        else:
+            # TODO: case cover by test
+            return str(value)
+    if isinstance(value, (bool, int, str, type(None))):
         return value
-    if isinstance(value, datetime.datetime):
-        return value.replace(microsecond=0).isoformat()
+    if isinstance(value, datetime.date):
+        if isinstance(value, datetime.datetime):
+            return value.replace(microsecond=0).isoformat()
+        else:
+            # TODO: cover case by test
+            return value.isoformat()
+
     try:
         if value.ndim == 0:
+            try:
+                # xarray.DataArray case
+                value = value.values
+            except AttributeError:
+                pass
+            if np.issubdtype(value.dtype, np.floating):
+                if np.isfinite(value):
+                    return float(value)
+                else:
+                    # TODO: cover case by test
+                    return str(value)
             if np.issubdtype(value.dtype, np.bool_):
                 return bool(value)
+            if np.issubdtype(value.dtype, np.str_):
+                return str(value)
             if np.issubdtype(value.dtype, np.integer):
                 return int(value)
-            if np.issubdtype(value.dtype, np.floating):
-                return float(value)
             if np.issubdtype(value.dtype, np.datetime64):
                 return np.datetime_as_string(value, unit="s")
+            # TODO: cover case by test
+            raise ValueError(
+                f"cannot serialize 0-d array of type {value.dtype}: {value}"
+            )
     except AttributeError:
         pass
-    # TODO: this is not right, handle remaining cases or raise
-    return value
+
+    if isinstance(value, dict):
+        # TODO: cover case by test
+        return {k: to_json(v) for k, v in value.items()}
+
+    return [to_json(v) for v in value]
 
 
 def get_dyn_config_attrs_env(ds: xr.Dataset, **kwargs):