Skip to content
This repository has been archived by the owner on Apr 26, 2024. It is now read-only.

Commit

Permalink
Eliminate a few Anys in LruCache type hints (#11453)
Browse files Browse the repository at this point in the history
  • Loading branch information
squahtx authored Nov 30, 2021
1 parent 432a174 commit 5a0b652
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 19 deletions.
1 change: 1 addition & 0 deletions changelog.d/11453.misc
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Improve type hints for `LruCache`.
9 changes: 8 additions & 1 deletion synapse/util/caches/deferred_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
Iterable,
MutableMapping,
Optional,
Sized,
TypeVar,
Union,
cast,
Expand Down Expand Up @@ -104,7 +105,13 @@ def metrics_cb() -> None:
max_size=max_entries,
cache_name=name,
cache_type=cache_type,
size_callback=(lambda d: len(d) or 1) if iterable else None,
size_callback=(
(lambda d: len(cast(Sized, d)) or 1)
# Argument 1 to "len" has incompatible type "VT"; expected "Sized"
# We trust that `VT` is `Sized` when `iterable` is `True`
if iterable
else None
),
metrics_collection_callback=metrics_cb,
apply_cache_factor_from_config=apply_cache_factor_from_config,
prune_unread_entries=prune_unread_entries,
Expand Down
37 changes: 21 additions & 16 deletions synapse/util/caches/lrucache.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,15 @@
import logging
import threading
import weakref
from enum import Enum
from functools import wraps
from typing import (
TYPE_CHECKING,
Any,
Callable,
Collection,
Dict,
Generic,
Iterable,
List,
Optional,
Type,
Expand Down Expand Up @@ -190,7 +191,7 @@ def __init__(
root: "ListNode[_Node]",
key: KT,
value: VT,
cache: "weakref.ReferenceType[LruCache]",
cache: "weakref.ReferenceType[LruCache[KT, VT]]",
clock: Clock,
callbacks: Collection[Callable[[], None]] = (),
prune_unread_entries: bool = True,
Expand Down Expand Up @@ -290,6 +291,12 @@ def move_to_front(self, clock: Clock, cache_list_root: ListNode) -> None:
self._global_list_node.update_last_access(clock)


class _Sentinel(Enum):
# defining a sentinel in this way allows mypy to correctly handle the
# type of a dictionary lookup.
sentinel = object()


class LruCache(Generic[KT, VT]):
"""
Least-recently-used cache, supporting prometheus metrics and invalidation callbacks.
Expand All @@ -302,7 +309,7 @@ def __init__(
max_size: int,
cache_name: Optional[str] = None,
cache_type: Type[Union[dict, TreeCache]] = dict,
size_callback: Optional[Callable] = None,
size_callback: Optional[Callable[[VT], int]] = None,
metrics_collection_callback: Optional[Callable[[], None]] = None,
apply_cache_factor_from_config: bool = True,
clock: Optional[Clock] = None,
Expand Down Expand Up @@ -339,7 +346,7 @@ def __init__(
else:
real_clock = clock

cache = cache_type()
cache: Union[Dict[KT, _Node[KT, VT]], TreeCache] = cache_type()
self.cache = cache # Used for introspection.
self.apply_cache_factor_from_config = apply_cache_factor_from_config

Expand Down Expand Up @@ -374,7 +381,7 @@ def __init__(
# creating more each time we create a `_Node`.
weak_ref_to_self = weakref.ref(self)

list_root = ListNode[_Node].create_root_node()
list_root = ListNode[_Node[KT, VT]].create_root_node()

lock = threading.Lock()

Expand Down Expand Up @@ -422,7 +429,7 @@ def cache_len() -> int:
def add_node(
key: KT, value: VT, callbacks: Collection[Callable[[], None]] = ()
) -> None:
node = _Node(
node: _Node[KT, VT] = _Node(
list_root,
key,
value,
Expand All @@ -439,10 +446,10 @@ def add_node(
if caches.TRACK_MEMORY_USAGE and metrics:
metrics.inc_memory_usage(node.memory)

def move_node_to_front(node: _Node) -> None:
def move_node_to_front(node: _Node[KT, VT]) -> None:
node.move_to_front(real_clock, list_root)

def delete_node(node: _Node) -> int:
def delete_node(node: _Node[KT, VT]) -> int:
node.drop_from_lists()

deleted_len = 1
Expand Down Expand Up @@ -496,7 +503,7 @@ def cache_get(

@synchronized
def cache_set(
key: KT, value: VT, callbacks: Iterable[Callable[[], None]] = ()
key: KT, value: VT, callbacks: Collection[Callable[[], None]] = ()
) -> None:
node = cache.get(key, None)
if node is not None:
Expand Down Expand Up @@ -590,8 +597,6 @@ def cache_clear() -> None:
def cache_contains(key: KT) -> bool:
return key in cache

self.sentinel = object()

# make sure that we clear out any excess entries after we get resized.
self._on_resize = evict

Expand All @@ -608,18 +613,18 @@ def cache_contains(key: KT) -> bool:
self.clear = cache_clear

def __getitem__(self, key: KT) -> VT:
result = self.get(key, self.sentinel)
if result is self.sentinel:
result = self.get(key, _Sentinel.sentinel)
if result is _Sentinel.sentinel:
raise KeyError()
else:
return cast(VT, result)
return result

def __setitem__(self, key: KT, value: VT) -> None:
self.set(key, value)

def __delitem__(self, key: KT, value: VT) -> None:
result = self.pop(key, self.sentinel)
if result is self.sentinel:
result = self.pop(key, _Sentinel.sentinel)
if result is _Sentinel.sentinel:
raise KeyError()

def __len__(self) -> int:
Expand Down
4 changes: 2 additions & 2 deletions synapse/util/linked_list.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ def remove_from_list(self) -> None:
# immediately rather than at the next GC.
self.cache_entry = None

def move_after(self, node: "ListNode") -> None:
def move_after(self, node: "ListNode[P]") -> None:
"""Move this node from its current location in the list to after the
given node.
"""
Expand Down Expand Up @@ -122,7 +122,7 @@ def _refs_remove_node_from_list(self) -> None:
self.prev_node = None
self.next_node = None

def _refs_insert_after(self, node: "ListNode") -> None:
def _refs_insert_after(self, node: "ListNode[P]") -> None:
"""Internal method to insert the node after the given node."""

# This method should only be called when we're not already in the list.
Expand Down

0 comments on commit 5a0b652

Please sign in to comment.