Add bounds to cache and clean key before caching

Author: eqbech

src/ert/gui/plotting/plot_api.py

@@ -3,8 +3,10 @@
 import io
 import json
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass
-from functools import cache, cached_property
+from functools import cached_property, lru_cache, wraps
+from inspect import signature
 from itertools import combinations as combi
 from typing import TYPE_CHECKING, Any, NamedTuple
 from urllib.parse import quote
@@ -58,6 +60,30 @@ class PlotApiKeyDefinition(NamedTuple):
     response: ResponseConfig | None = None
 
 
+def clear_suffix[**P, R](
+    key: str,
+    split: str
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Return a decorator that strips a trailing ``@<suffix>`` from the
+    argument named ``key`` before calling the wrapped function.
+    """
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        sig = signature(func)
+
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            bound = sig.bind(*args, **kwargs)
+            value = bound.arguments.get(key)
+            if isinstance(value, str) and "@" in value:
+                bound.arguments[key] = value.split(split, maxsplit=1)[0]
+            return func(*bound.args, **bound.kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 class PlotApi:
     def __init__(self, ens_path: Path) -> None:
         self.ens_path: Path = ens_path
@@ -344,10 +370,9 @@ def data_for_response(
                 return df
 
     @staticmethod
-    @cache
+    @lru_cache(maxsize=64)
+    @clear_suffix(key="key", split="@")
     def data_for_gradient(ensemble_id: str, key: str, ens_path: Path) -> pd.DataFrame:
-        if "@" in key:
-            key = key.split("@", maxsplit=1)[0]
         with create_ertserver_client(ens_path) as client:
             http_response = client.get(
                 f"/ensembles/{ensemble_id}/gradients/{PlotApi.escape(key)}",
@@ -371,7 +396,7 @@ def data_for_gradient(ensemble_id: str, key: str, ens_path: Path) -> pd.DataFram
             )
 
     @staticmethod
-    @cache
+    @lru_cache(maxsize=32)
     def data_for_controls(
         ensemble_id: str, parameter_keys: tuple[str, ...], ens_path: Path
     ) -> pd.DataFrame:
@@ -395,7 +420,7 @@ def data_for_controls(
         return pd.concat(frames, ignore_index=True)
 
     @staticmethod
-    @cache
+    @lru_cache(maxsize=64)
     def data_for_parameter(
         ensemble_id: str, parameter_key: str, ens_path: Path
     ) -> pd.DataFrame:

tests/ert/unit_tests/gui/tools/plot/test_plot_api.py

@@ -694,9 +694,9 @@ def test_that_data_for_gradient_is_fetched_once_for_repeated_calls(api_and_stora
     PlotApi.data_for_gradient.cache_clear()
 
     for _ in range(5):
-        _ = PlotApi.data_for_gradient(str(ensemble.id), objective_key, ens_path)
+        _ = PlotApi.data_for_gradient(str(ensemble.id), objective_key + "@filler", ens_path)
 
     # hits, misses, maxsize, currsize
-    expected_cache_info = (4, 1, None, 1)
+    expected_cache_info = (4, 1, 64, 1)
 
     assert PlotApi.data_for_gradient.cache_info() == expected_cache_info