isolate logic to cache fields in FieldCache

Author: amorison

src/stagpy/_caching.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import typing
+from collections import deque
+from dataclasses import dataclass
+from functools import cached_property
+
+if typing.TYPE_CHECKING:
+    from .datatypes import Field
+
+
+@dataclass(frozen=True)
+class FieldCache:
+    """FIFO cache of [Field][]s.
+
+    If `maxsize` is None, entries are never evicted from the cache.
+    """
+
+    maxsize: int | None
+
+    @cached_property
+    def _stack(self) -> deque[tuple[int, str]]:
+        return deque()
+
+    @cached_property
+    def _data(self) -> dict[tuple[int, str], Field]:
+        return {}
+
+    def _prune(self) -> None:
+        if self.maxsize is None:
+            return
+        while len(self._stack) > self.maxsize:
+            elt = self._stack.popleft()
+            del self._data[elt]
+
+    def resize(self, new_size: int | None) -> None:
+        object.__setattr__(self, "maxsize", new_size)
+        self._prune()
+
+    def insert(self, istep: int, name: str, field: Field) -> None:
+        key = (istep, name)
+        if key not in self._data:
+            self._stack.append(key)
+        self._data[key] = field
+        self._prune()
+
+    def get(self, istep: int, name: str) -> Field | None:
+        return self._data.get((istep, name))
+
+    def evict_istep(self, istep: int) -> None:
+        to_keep = []
+        for key in self._stack:
+            if key[0] == istep:
+                del self._data[key]
+            else:
+                to_keep.append(key)
+        self._stack.clear()
+        self._stack.extend(to_keep)
+        assert len(self._stack) == len(self._data)

src/stagpy/stagyydata.py

@@ -20,6 +20,7 @@
 
 from . import _helpers, error, phyvars, stagyyparsers, step
 from . import datatypes as dt
+from ._caching import FieldCache
 from .parfile import StagyyPar
 from .stagyyparsers import FieldXmf, TracersXmf
 from .step import Step
@@ -332,9 +333,7 @@ def __getitem__(self, istep: int | slice | Sequence[StepIndex]) -> Step | StepsV
 
     def __delitem__(self, istep: int | None) -> None:
         if istep is not None and istep in self._data:
-            self.sdat._collected_fields = [
-                (i, f) for i, f in self.sdat._collected_fields if i != istep
-            ]
+            self.sdat._field_cache.evict_istep(istep)
             del self._data[istep]
 
     def __len__(self) -> int:
@@ -668,9 +667,6 @@ def __init__(self, path: PathLike, read_parameters_dat: bool = True):
         self.steps = Steps(self)
         self.snaps = Snaps(self)
         self._read_parameters_dat = read_parameters_dat
-        self._nfields_max: int | None = 50
-        # list of (istep, field_name) in memory
-        self._collected_fields: list[tuple[int, str]] = []
 
     def __repr__(self) -> str:
         return f"StagyyData({self.path!r})"
@@ -752,23 +748,17 @@ def _files(self) -> set[Path]:
             return set(out_dir.iterdir())
         return set()
 
-    @property
-    def nfields_max(self) -> int | None:
-        """Maximum number of scalar fields kept in memory.
+    def set_nfields_max(self, nfields: int | None) -> None:
+        """Adjust maximum number of scalar fields kept in memory.
 
         Setting this to a value lower or equal to 5 raises a
         [stagpy.error.InvalidNfieldsError][].  Set this to `None` if
         you do not want any limit on the number of scalar fields kept in
         memory.  Defaults to 50.
         """
-        return self._nfields_max
-
-    @nfields_max.setter
-    def nfields_max(self, nfields: int | None) -> None:
-        """Check nfields > 5 or None."""
         if nfields is not None and nfields <= 5:
             raise error.InvalidNfieldsError(nfields)
-        self._nfields_max = nfields
+        self._field_cache.resize(nfields)
 
     def filename(
         self,
@@ -811,3 +801,7 @@ def _binfiles_set(self, isnap: int) -> set[Path]:
             for fstem in phyvars.FIELD_FILES
         )
         return possible_files & self._files
+
+    @cached_property
+    def _field_cache(self) -> FieldCache:
+        return FieldCache(maxsize=50)

src/stagpy/step.py

@@ -25,6 +25,7 @@
     from numpy.typing import NDArray
     from pandas import DataFrame, Series
 
+    from ._caching import FieldCache
     from .datatypes import Varf
     from .stagyydata import StagyyData
 
@@ -292,28 +293,41 @@ def __init__(
         self._extra = extravars
         self._files = files
         self._filesh5 = filesh5
-        self._data: dict[str, Field] = {}
         super().__init__()
 
+    @cached_property
+    def _all_vars(self) -> set[str]:
+        return set(self._vars.keys()).union(self._extra.keys())
+
+    @cached_property
+    def _cache(self) -> FieldCache:
+        return self.step.sdat._field_cache
+
     def __getitem__(self, name: str) -> Field:
-        if name in self._data:
-            return self._data[name]
-        if name in self._vars:
-            fld_names, parsed_data = self._get_raw_data(name)
-        elif name in self._extra:
-            self._data[name] = self._extra[name](self.step)
-            return self._data[name]
-        else:
+        if name not in self._all_vars:
             raise error.UnknownFieldVarError(name)
+
+        maybe_fld = self._cache.get(self.step.istep, name)
+        if maybe_fld is not None:
+            return maybe_fld
+
+        if name in self._extra:
+            fld = self._extra[name](self.step)
+            self._cache.insert(self.step.istep, name, fld)
+            return fld
+
+        # requested field is one of self._vars
+        fld_names, parsed_data = self._get_raw_data(name)
         if parsed_data is None:
             raise error.MissingDataError(
                 f"Missing field {name} in step {self.step.istep}"
             )
         header, fields = parsed_data
         self._cropped__header = header
-        for fld_name, fld in zip(fld_names, fields):
-            self._set(fld_name, fld)
-        return self._data[name]
+        for fld_name, fld_vals in zip(fld_names, fields):
+            fld = Field(fld_vals, self._vars[fld_name])
+            self._cache.insert(self.step.istep, fld_name, fld)
+        return self[name]
 
     @cached_property
     def _present_fields(self) -> list[str]:
@@ -371,20 +385,6 @@ def _get_raw_data(self, name: str) -> tuple[list[str], Any]:
                     break
         return list_fvar, parsed_data
 
-    def _set(self, name: str, fld: NDArray) -> None:
-        sdat = self.step.sdat
-        col_fld = sdat._collected_fields
-        col_fld.append((self.step.istep, name))
-        if sdat.nfields_max is not None:
-            while len(col_fld) > sdat.nfields_max:
-                istep, fld_name = col_fld.pop(0)
-                del sdat.steps[istep].fields[fld_name]
-        self._data[name] = Field(fld, self._vars[name])
-
-    def __delitem__(self, name: str) -> None:
-        if name in self._data:
-            del self._data[name]
-
     @cached_property
     def _header(self) -> dict[str, Any] | None:
         if self.step.isnap is None:

tests/test_stagyydata.py

@@ -11,23 +11,9 @@ def test_sdat_path(example_dir: Path, sdat: StagyyData) -> None:
     assert sdat.path == example_dir
 
 
-def test_sdat_deflt_nfields_max(sdat: StagyyData) -> None:
-    assert sdat.nfields_max == 50
-
-
-def test_sdat_set_nfields_max(sdat: StagyyData) -> None:
-    sdat.nfields_max = 6
-    assert sdat.nfields_max == 6
-
-
-def test_sdat_set_nfields_max_none(sdat: StagyyData) -> None:
-    sdat.nfields_max = None
-    assert sdat.nfields_max is None
-
-
 def test_sdat_set_nfields_max_invalid(sdat: StagyyData) -> None:
     with pytest.raises(stagpy.error.InvalidNfieldsError) as err:
-        sdat.nfields_max = 5
+        sdat.set_nfields_max(5)
     assert err.value.nfields == 5