Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Annotate LRUCache #3395

Merged
merged 1 commit into from
Oct 12, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions xarray/backends/file_manager.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,18 @@
import contextlib
import io
import threading
import warnings
from typing import Any, Dict
from typing import Any, Dict, cast

from ..core import utils
from ..core.options import OPTIONS
from .locks import acquire
from .lru_cache import LRUCache

# Global cache for storing open files.
FILE_CACHE = LRUCache(OPTIONS["file_cache_maxsize"], on_evict=lambda k, v: v.close())
FILE_CACHE: LRUCache[str, io.IOBase] = LRUCache(
maxsize=cast(int, OPTIONS["file_cache_maxsize"]), on_evict=lambda k, v: v.close()
)
assert FILE_CACHE.maxsize, "file cache must be at least size one"


Expand Down
41 changes: 26 additions & 15 deletions xarray/backends/lru_cache.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
import collections
import collections.abc
import threading
from collections import OrderedDict
from typing import Any, Callable, Iterator, MutableMapping, Optional, TypeVar

K = TypeVar("K")
V = TypeVar("V")

class LRUCache(collections.abc.MutableMapping):

class LRUCache(MutableMapping[K, V]):
"""Thread-safe LRUCache based on an OrderedDict.

All dict operations (__getitem__, __setitem__, __contains__) update the
Expand All @@ -18,7 +21,14 @@ class LRUCache(collections.abc.MutableMapping):
the cache, e.g., ``cache.maxsize = new_size``.
"""

def __init__(self, maxsize, on_evict=None):
_cache: "OrderedDict[K, V]"
_maxsize: int
_lock: threading.RLock
_on_evict: Optional[Callable[[K, V], Any]]

__slots__ = ("_cache", "_lock", "_maxsize", "_on_evict")

def __init__(self, maxsize: int, on_evict: Callable[[K, V], Any] = None):
"""
Parameters
----------
Expand All @@ -33,25 +43,26 @@ def __init__(self, maxsize, on_evict=None):
if maxsize < 0:
raise ValueError("maxsize must be non-negative")
self._maxsize = maxsize
self._on_evict = on_evict
self._cache = collections.OrderedDict()
self._cache = OrderedDict()
self._lock = threading.RLock()
self._on_evict = on_evict

def __getitem__(self, key):
def __getitem__(self, key: K) -> V:
# record recent use of the key by moving it to the front of the list
with self._lock:
value = self._cache[key]
self._cache.move_to_end(key)
return value

def _enforce_size_limit(self, capacity):
"""Shrink the cache if necessary, evicting the oldest items."""
def _enforce_size_limit(self, capacity: int) -> None:
"""Shrink the cache if necessary, evicting the oldest items.
"""
while len(self._cache) > capacity:
key, value = self._cache.popitem(last=False)
if self._on_evict is not None:
self._on_evict(key, value)

def __setitem__(self, key, value):
def __setitem__(self, key: K, value: V) -> None:
with self._lock:
if key in self._cache:
# insert the new value at the end
Expand All @@ -65,24 +76,24 @@ def __setitem__(self, key, value):
# not saving, immediately evict
self._on_evict(key, value)

def __delitem__(self, key):
def __delitem__(self, key: K) -> None:
del self._cache[key]

def __iter__(self):
def __iter__(self) -> Iterator[K]:
# create a list, so accessing the cache during iteration cannot change
# the iteration order
return iter(list(self._cache))

def __len__(self):
def __len__(self) -> int:
return len(self._cache)

@property
def maxsize(self):
def maxsize(self) -> int:
"""Maximum number of items can be held in the cache."""
return self._maxsize

@maxsize.setter
def maxsize(self, size):
def maxsize(self, size: int) -> None:
"""Resize the cache, evicting the oldest items if necessary."""
if size < 0:
raise ValueError("maxsize must be non-negative")
Expand Down