Skip to content

Commit

Permalink
Update test and nicer code
Browse files Browse the repository at this point in the history
  • Loading branch information
bouweandela committed Dec 10, 2024
1 parent eb0b19f commit 916a1df
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 13 deletions.
22 changes: 10 additions & 12 deletions lib/iris/_lazy_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

from functools import lru_cache, wraps
from types import ModuleType
from typing import Any, Sequence
from typing import Sequence

import dask
import dask.array as da
Expand Down Expand Up @@ -229,24 +229,24 @@ def _optimum_chunksize(


class LRUCache:
def __init__(self, maxsize: int):
self._cache: dict[str, Any] = {}
def __init__(self, maxsize: int) -> None:
self._cache: dict[str, da.Array] = {}
self.maxsize = maxsize

def __getitem__(self, key):
def __getitem__(self, key: str) -> da.Array:
value = self._cache.pop(key)
self._cache[key] = value
return value

def __setitem__(self, key, value):
def __setitem__(self, key: str, value: da.Array) -> None:
self._cache[key] = value
if len(self._cache) > self.maxsize:
self._cache.pop(next(iter(self._cache)))

def __contains__(self, key):
def __contains__(self, key: str) -> bool:
return key in self._cache

def __repr__(self):
def __repr__(self) -> str:
return (
f"<{self.__class__.__name__} maxsize={self.maxsize} cache={self._cache!r} >"
)
Expand Down Expand Up @@ -331,12 +331,10 @@ def as_lazy_data(data, chunks=None, asarray=False, meta=None, dims_fixed=None):

if is_lazy_data(data):
result = data
elif key in CACHE:
result = CACHE[key].copy()
else:
result = da.from_array(data, chunks=chunks, asarray=asarray, meta=meta)
if key is not None:
CACHE[key] = result.copy()
if key not in CACHE:
CACHE[key] = da.from_array(data, chunks=chunks, asarray=asarray, meta=meta)
result = CACHE[key].copy()

return result

Expand Down
2 changes: 1 addition & 1 deletion lib/iris/tests/unit/lazy_data/test_as_lazy_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ def test_default_chunks_limiting(self, mocker):
as_lazy_data(data)
assert limitcall_patch.call_args_list == [
mock.call(
list(test_shape),
tuple(test_shape),
shape=test_shape,
dtype=np.dtype("f4"),
dims_fixed=None,
Expand Down

0 comments on commit 916a1df

Please sign in to comment.