State: quick fix to v0 (NOAA-GFDL#245)

* Use `tmpdir` of `pytest` for a cleaner utest

* Allow `None` as a valid State attribute

* Deal with `None` iinput in move & copy
Enhance errors on copy

Author: FlorianDeconinck

ndsl/quantity/state.py

@@ -2,7 +2,7 @@
 
 import dataclasses
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Self
+from typing import TYPE_CHECKING, Any, Callable, Self, TypeAlias
 
 import dacite
 import xarray as xr
@@ -13,6 +13,8 @@
 if TYPE_CHECKING:
     from ndsl import QuantityFactory
 
+StateMemoryMapping: TypeAlias = dict[str, dict | ArrayLike | None]
+
 
 @dataclasses.dataclass
 class State:
@@ -86,7 +88,7 @@ def ones(cls, quantity_factory: QuantityFactory) -> Self:
     def copy_memory(
         cls,
         quantity_factory: QuantityFactory,
-        memory_map: dict[str, Any],
+        memory_map: StateMemoryMapping,
     ) -> Self:
         """Allocate all quantities and fill their value based
         on the given memory map. See `update_from_memory`"""
@@ -100,7 +102,7 @@ def copy_memory(
     def move_memory(
         cls,
         quantity_factory: QuantityFactory,
-        memory_map: dict[str, Any],
+        memory_map: StateMemoryMapping,
         *,
         check_shape_and_strides: bool = True,
     ) -> Self:
@@ -147,10 +149,14 @@ class InnerA
 
         def _update_from_memory_recursive(
             state: State,
-            memory_map: dict[str, dict | ArrayLike],
+            memory_map: StateMemoryMapping,
         ):
             for name, array in memory_map.items():
-                if isinstance(array, dict):
+                if array is None:
+                    raise TypeError(
+                        f"State memory copy: illegal copy from None for attribute {name}"
+                    )
+                elif isinstance(array, dict):
                     _update_from_memory_recursive(state.__getattribute__(name), array)
                 else:
                     try:
@@ -165,7 +171,7 @@ def _update_from_memory_recursive(
 
     def update_move_memory(
         self,
-        memory_map: dict[str, dict | ArrayLike],
+        memory_map: StateMemoryMapping,
         *,
         check_shape_and_strides: bool = True,
     ) -> None:
@@ -204,11 +210,11 @@ class InnerA
                 shape and strides as the original quantity
         """
 
-        def _update_zero_copy_recursive(
-            state: State, memory_map: dict[str, dict | ArrayLike]
-        ):
+        def _update_zero_copy_recursive(state: State, memory_map: StateMemoryMapping):
             for name, array in memory_map.items():
-                if isinstance(array, dict):
+                if array is None:
+                    state.__setattr__(name, None)
+                elif isinstance(array, dict):
                     _update_zero_copy_recursive(state.__getattribute__(name), array)
                 else:
                     quantity = state.__getattribute__(name)
@@ -227,8 +233,11 @@ def _update_zero_copy_recursive(
                                 f"  Strides: {array.strides} != {quantity.data.strides}"
                             )
                             raise e
-
-                    quantity.data = array
+                    try:
+                        quantity.data = array
+                    except Exception as e:
+                        e.add_note(f"  Error on {name} for {type(state)}")
+                        raise e
 
         _update_zero_copy_recursive(self, memory_map)
 

tests/quantity/test_state.py

@@ -47,16 +47,16 @@ class InnerB:
     )
 
 
-def test_state():
+def test_state(tmpdir):
     _, quantity_factory = get_factories_single_tile(
         5, 5, 3, 0, backend="dace:cpu_kfirst"
     )
 
     microphys_state = CodeState.zeros(quantity_factory)
     microphys_state.inner_A.A.field[:] = 42.42
-    microphys_state.to_netcdf()
+    microphys_state.to_netcdf(Path(tmpdir))
     microphys_state2 = CodeState.zeros(quantity_factory)
-    microphys_state2.update_from_netcdf(Path("./"))
+    microphys_state2.update_from_netcdf(Path(tmpdir))
     assert (microphys_state2.inner_A.A.field[:] == 42.42).all()
     a = np.ones((5, 5, 3))
     b = np.ones((5, 5, 3))