Merge pull request #52 from zwicker-group/types

Cleaned up some types

modelrunner/storage/backend/zarr.py

@@ -226,7 +226,7 @@ def _read_object(self, loc: Sequence[str]) -> Any:
         return self[loc][0]
 
     def _write_object(self, loc: Sequence[str], obj: Any) -> None:
-        arr = np.empty(1, dtype=object)  # encode object in an array
+        arr: np.ndarray = np.empty(1, dtype=object)  # encode object in an array
         arr[0] = obj
         parent, name = self._get_parent(loc)
         parent.array(name, arr, object_codec=self.codec, overwrite=True)

modelrunner/storage/trajectory.py

@@ -13,7 +13,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Dict, Iterator, Literal, Optional
+from typing import Any, Dict, Iterator, Literal, Optional, Tuple
 
 import numpy as np
 
@@ -114,7 +114,7 @@ def append(self, data: Any, time: Optional[float] = None) -> None:
             # initialize new trajectory
             if isinstance(data, np.ndarray):
                 dtype = data.dtype
-                shape = data.shape
+                shape: Tuple[int, ...] = data.shape
                 self._item_type = "array"
             else:
                 dtype = object
@@ -132,7 +132,7 @@ def append(self, data: Any, time: Optional[float] = None) -> None:
         if self._item_type == "array":
             self._trajectory.extend_dynamic_array("data", data)
         elif self._item_type == "object":
-            arr = np.empty((), dtype=object)
+            arr: np.ndarray = np.empty((), dtype=object)
             arr[...] = data
             self._trajectory.extend_dynamic_array("data", arr)
         else:

tests/test_results.py

@@ -127,7 +127,7 @@ def test_result_collections():
     assert rc1 == ResultCollection([r1, r2, r3])
 
     # test result dataframes
-    rc1.dataframe
+    rc1.as_dataframe()
 
 
 def test_collection_groupby():