Merge remote-tracking branch 'origin/forman-60-custom_attrs' into forman-60-custom_attrs

Author: forman

tests/config/test_attrs.py

@@ -2,11 +2,12 @@
 # Permissions are hereby granted under the terms of the MIT License:
 # https://opensource.org/licenses/MIT.
 
+import datetime
 import json
 import unittest
-from datetime import datetime
 
 import numpy as np
+import xarray as xr
 import pytest
 
 from zappend.config.attrs import ConfigAttrsUserFunctions
@@ -148,52 +149,108 @@ def test_yes(self):
 
 
 class EvalExprTest(unittest.TestCase):
-    def test_all_cases(self):
+    def test_scalar_result(self):
+        # scalars
+        self.assertEqual(None, eval_expr("None", {}))
         self.assertEqual(True, eval_expr("True", {}))
         self.assertEqual(13, eval_expr("13", {}))
+        self.assertEqual("nan", eval_expr("float('NaN')", {}))
         self.assertEqual(0.5, eval_expr("0.5", {}))
         self.assertEqual("ABC", eval_expr("'ABC'", {}))
-
-        switches = np.array([True, False])
+        time = datetime.date.fromisoformat("2024-01-02")
         self.assertEqual(
-            True,
-            eval_expr("switches[0]", {"switches": switches}),
+            "2024-01-02",
+            eval_expr("time", dict(time=time)),
         )
-
-        levels = [3, 4, 5]
-        self.assertIs(
-            levels,
-            eval_expr("levels", {"levels": levels}),
+        time = datetime.datetime.fromisoformat("2024-01-02T10:20:30")
+        self.assertEqual(
+            "2024-01-02T10:20:30",
+            eval_expr("time", dict(time=time)),
         )
+        with pytest.raises(
+            ValueError, match="cannot serialize value of type <class 'object'>"
+        ):
+            eval_expr("obj", dict(obj=object()))
 
-        levels = np.array([3, 4, 5])
+    def test_dict_result(self):
+        self.assertEqual({}, eval_expr("d", dict(d={})))
         self.assertEqual(
-            3,
-            eval_expr("levels[0]", {"levels": levels}),
+            {
+                "b": True,
+                "i": 13,
+                "t": [1, "B", {}],
+                "f": 13.2,
+                "d": "2020-05-04",
+                "np_a": [0.1, 0.2],
+                "xr_a": [0.3, 0.4],
+            },
+            eval_expr(
+                "d",
+                dict(
+                    d={
+                        "b": True,
+                        "i": 13,
+                        "t": (1, "B", {}),
+                        "f": 13.2,
+                        "d": datetime.date.fromisoformat("2020-05-04"),
+                        "np_a": np.array([0.1, 0.2]),
+                        "xr_a": xr.DataArray(np.array([0.3, 0.4])),
+                    }
+                ),
+            ),
         )
 
-        lon = np.array([11.05, 11.15, 11.25])
-        self.assertIs(
-            lon,
-            eval_expr("lon", {"lon": lon}),
-        )
+    def test_array_1d_result(self):
+        # 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([11.05, "nan", 11.25], dtype="float64")
+        self.assert_array_ok(["A", "B"])
+        self.assert_array_ok(["2024-01-02T10:20:30"], dtype="datetime64[s]")
+        with pytest.raises(
+            ValueError,
+            match=(
+                "cannot serialize 0-d array"
+                " of type <class 'numpy.ndarray'>, dtype=dtype\\('O'\\)"
+            ),
+        ):
+            eval_expr("a", dict(a=xr.DataArray([object(), object()])))
 
-        lon = np.array([11.05, 11.15, 11.25])
+    def test_array_2d_result(self):
+        self.assert_array_ok([[3, 4], [5, 6]])
+
+    def assert_array_ok(self, a: list, dtype=None):
+        # Test list
+        self.assertEqual(
+            a,
+            eval_expr("a", dict(a=a)),
+        )
         self.assertEqual(
-            11.05,
-            eval_expr("lon[0]", {"lon": lon}),
+            a[0],
+            eval_expr("a[0]", dict(a=a)),
         )
 
-        names = np.array(["A", "B"])
+        # Test numpy.ndarray
+        np_a = np.array(a, dtype=dtype) if dtype is not None else np.array(a)
+        self.assertEqual(
+            a,
+            eval_expr("a", dict(a=np_a)),
+        )
         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)
         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)),
         )