[mypy] Annotate other/lru_cache and other/lfu_cache (#5755)

* Adds repr and doctest of current behavior linkedlist in other/lru_cache * Blocks removal of head or tail of double linked list * clarifies add() logic for double linked list in other/lru_cache * expands doctests to compare cache and lru cache * [mypy] annotates vars for other/lru_cache * [mypy] Annotates lru_cache decorator for other/lru_cache * Higher order functions require a verbose Callable annotation * [mypy] Makes LRU_Cache generic over key and value types for other/lru_cache + no reason to force int -> int * [mypy] makes decorator a classmethod for access to class generic types * breaks two long lines in doctest for other/lru_cache * simplifies boundary test remove() for other/lru_cache * [mypy] Annotates, adds doctests, and makes Generic other/lfu_cache See also commits to other/lru_cache which guided these * [mypy] annotates cls var in other/lfu_cache * cleans up items from code review for lfu_cache and lru_cache * [mypy] runs mypy on lfu_cache and lru_cache

Author: spazm

mypy.ini

@@ -2,4 +2,4 @@
 ignore_missing_imports = True
 install_types = True
 non_interactive = True
-exclude = (other/least_recently_used.py|other/lfu_cache.py|other/lru_cache.py)
+exclude = (other/least_recently_used.py)

other/lfu_cache.py

@@ -1,61 +1,165 @@
 from __future__ import annotations
 
-from typing import Callable
+from typing import Callable, Generic, TypeVar
 
+T = TypeVar("T")
+U = TypeVar("U")
 
-class DoubleLinkedListNode:
+
+class DoubleLinkedListNode(Generic[T, U]):
  """
  Double Linked List Node built specifically for LFU Cache
+
+ >>> node = DoubleLinkedListNode(1,1)
+ >>> node
+ Node: key: 1, val: 1, freq: 0, has next: False, has prev: False
  """
 
- def __init__(self, key: int, val: int):
+ def __init__(self, key: T | None, val: U | None):
  self.key = key
  self.val = val
- self.freq = 0
- self.next = None
- self.prev = None
+ self.freq: int = 0
+ self.next: DoubleLinkedListNode[T, U] | None = None
+ self.prev: DoubleLinkedListNode[T, U] | None = None
 
+ def __repr__(self) -> str:
+ return "Node: key: {}, val: {}, freq: {}, has next: {}, has prev: {}".format(
+ self.key, self.val, self.freq, self.next is not None, self.prev is not None
+ )
 
-class DoubleLinkedList:
+
+class DoubleLinkedList(Generic[T, U]):
  """
  Double Linked List built specifically for LFU Cache
+
+ >>> dll: DoubleLinkedList = DoubleLinkedList()
+ >>> dll
+ DoubleLinkedList,
+ Node: key: None, val: None, freq: 0, has next: True, has prev: False,
+ Node: key: None, val: None, freq: 0, has next: False, has prev: True
+
+ >>> first_node = DoubleLinkedListNode(1,10)
+ >>> first_node
+ Node: key: 1, val: 10, freq: 0, has next: False, has prev: False
+
+
+ >>> dll.add(first_node)
+ >>> dll
+ DoubleLinkedList,
+ Node: key: None, val: None, freq: 0, has next: True, has prev: False,
+ Node: key: 1, val: 10, freq: 1, has next: True, has prev: True,
+ Node: key: None, val: None, freq: 0, has next: False, has prev: True
+
+ >>> # node is mutated
+ >>> first_node
+ Node: key: 1, val: 10, freq: 1, has next: True, has prev: True
+
+ >>> second_node = DoubleLinkedListNode(2,20)
+ >>> second_node
+ Node: key: 2, val: 20, freq: 0, has next: False, has prev: False
+
+ >>> dll.add(second_node)
+ >>> dll
+ DoubleLinkedList,
+ Node: key: None, val: None, freq: 0, has next: True, has prev: False,
+ Node: key: 1, val: 10, freq: 1, has next: True, has prev: True,
+ Node: key: 2, val: 20, freq: 1, has next: True, has prev: True,
+ Node: key: None, val: None, freq: 0, has next: False, has prev: True
+
+ >>> removed_node = dll.remove(first_node)
+ >>> assert removed_node == first_node
+ >>> dll
+ DoubleLinkedList,
+ Node: key: None, val: None, freq: 0, has next: True, has prev: False,
+ Node: key: 2, val: 20, freq: 1, has next: True, has prev: True,
+ Node: key: None, val: None, freq: 0, has next: False, has prev: True
+
+
+ >>> # Attempt to remove node not on list
+ >>> removed_node = dll.remove(first_node)
+ >>> removed_node is None
+ True
+
+ >>> # Attempt to remove head or rear
+ >>> dll.head
+ Node: key: None, val: None, freq: 0, has next: True, has prev: False
+ >>> dll.remove(dll.head) is None
+ True
+
+ >>> # Attempt to remove head or rear
+ >>> dll.rear
+ Node: key: None, val: None, freq: 0, has next: False, has prev: True
+ >>> dll.remove(dll.rear) is None
+ True
+
+
  """
 
- def __init__(self):
- self.head = DoubleLinkedListNode(None, None)
- self.rear = DoubleLinkedListNode(None, None)
+ def __init__(self) -> None:
+ self.head: DoubleLinkedListNode[T, U] = DoubleLinkedListNode(None, None)
+ self.rear: DoubleLinkedListNode[T, U] = DoubleLinkedListNode(None, None)
  self.head.next, self.rear.prev = self.rear, self.head
 
- def add(self, node: DoubleLinkedListNode) -> None:
+ def __repr__(self) -> str:
+ rep = ["DoubleLinkedList"]
+ node = self.head
+ while node.next is not None:
+ rep.append(str(node))
+ node = node.next
+ rep.append(str(self.rear))
+ return ",\n ".join(rep)
+
+ def add(self, node: DoubleLinkedListNode[T, U]) -> None:
  """
- Adds the given node at the head of the list and shifting it to proper position
+ Adds the given node at the tail of the list and shifting it to proper position
  """
 
- temp = self.rear.prev
+ previous = self.rear.prev
+
+ # All nodes other than self.head are guaranteed to have non-None previous
+ assert previous is not None
 
- self.rear.prev, node.next = node, self.rear
- temp.next, node.prev = node, temp
+ previous.next = node
+ node.prev = previous
+ self.rear.prev = node
+ node.next = self.rear
  node.freq += 1
  self._position_node(node)
 
- def _position_node(self, node: DoubleLinkedListNode) -> None:
- while node.prev.key and node.prev.freq > node.freq:
- node1, node2 = node, node.prev
- node1.prev, node2.next = node2.prev, node1.prev
- node1.next, node2.prev = node2, node1
+ def _position_node(self, node: DoubleLinkedListNode[T, U]) -> None:
+ """
+ Moves node forward to maintain invariant of sort by freq value
+ """
+
+ while node.prev is not None and node.prev.freq > node.freq:
+ # swap node with previous node
+ previous_node = node.prev
 
- def remove(self, node: DoubleLinkedListNode) -> DoubleLinkedListNode:
+ node.prev = previous_node.prev
+ previous_node.next = node.prev
+ node.next = previous_node
+ previous_node.prev = node
+
+ def remove(
+ self, node: DoubleLinkedListNode[T, U]
+ ) -> DoubleLinkedListNode[T, U] | None:
  """
  Removes and returns the given node from the list
+
+ Returns None if node.prev or node.next is None
  """
 
- temp_last, temp_next = node.prev, node.next
- node.prev, node.next = None, None
- temp_last.next, temp_next.prev = temp_next, temp_last
+ if node.prev is None or node.next is None:
+ return None
+
+ node.prev.next = node.next
+ node.next.prev = node.prev
+ node.prev = None
+ node.next = None
  return node
 
 
-class LFUCache:
+class LFUCache(Generic[T, U]):
  """
  LFU Cache to store a given capacity of data. Can be used as a stand-alone object
  or as a function decorator.
@@ -66,9 +170,11 @@ class LFUCache:
  >>> cache.get(1)
  1
  >>> cache.set(3, 3)
- >>> cache.get(2) # None is returned
+ >>> cache.get(2) is None
+ True
  >>> cache.set(4, 4)
- >>> cache.get(1) # None is returned
+ >>> cache.get(1) is None
+ True
  >>> cache.get(3)
  3
  >>> cache.get(4)
@@ -89,15 +195,15 @@ class LFUCache:
  """
 
  # class variable to map the decorator functions to their respective instance
- decorator_function_to_instance_map = {}
+ decorator_function_to_instance_map: dict[Callable[[T], U], LFUCache[T, U]] = {}
 
  def __init__(self, capacity: int):
- self.list = DoubleLinkedList()
+ self.list: DoubleLinkedList[T, U] = DoubleLinkedList()
  self.capacity = capacity
  self.num_keys = 0
  self.hits = 0
  self.miss = 0
- self.cache = {}
+ self.cache: dict[T, DoubleLinkedListNode[T, U]] = {}
 
  def __repr__(self) -> str:
  """
@@ -110,73 +216,93 @@ def __repr__(self) -> str:
  f"capacity={self.capacity}, current_size={self.num_keys})"
  )
 
- def __contains__(self, key: int) -> bool:
+ def __contains__(self, key: T) -> bool:
  """
  >>> cache = LFUCache(1)
+
  >>> 1 in cache
  False
+
  >>> cache.set(1, 1)
  >>> 1 in cache
  True
  """
+
  return key in self.cache
 
- def get(self, key: int) -> int | None:
+ def get(self, key: T) -> U | None:
  """
  Returns the value for the input key and updates the Double Linked List. Returns
- None if key is not present in cache
+ Returns None if key is not present in cache
  """
 
  if key in self.cache:
  self.hits += 1
- self.list.add(self.list.remove(self.cache[key]))
- return self.cache[key].val
+ value_node: DoubleLinkedListNode[T, U] = self.cache[key]
+ node = self.list.remove(self.cache[key])
+ assert node == value_node
+
+ # node is guaranteed not None because it is in self.cache
+ assert node is not None
+ self.list.add(node)
+ return node.val
  self.miss += 1
  return None
 
- def set(self, key: int, value: int) -> None:
+ def set(self, key: T, value: U) -> None:
  """
  Sets the value for the input key and updates the Double Linked List
  """
 
  if key not in self.cache:
  if self.num_keys >= self.capacity:
- key_to_delete = self.list.head.next.key
- self.list.remove(self.cache[key_to_delete])
- del self.cache[key_to_delete]
+ # delete first node when over capacity
+ first_node = self.list.head.next
+
+ # guaranteed to have a non-None first node when num_keys > 0
+ # explain to type checker via assertions
+ assert first_node is not None
+ assert first_node.key is not None
+ assert self.list.remove(first_node) is not None
+ # first_node guaranteed to be in list
+
+ del self.cache[first_node.key]
  self.num_keys -= 1
  self.cache[key] = DoubleLinkedListNode(key, value)
  self.list.add(self.cache[key])
  self.num_keys += 1
 
  else:
  node = self.list.remove(self.cache[key])
+ assert node is not None # node guaranteed to be in list
  node.val = value
  self.list.add(node)
 
- @staticmethod
- def decorator(size: int = 128):
+ @classmethod
+ def decorator(
+ cls: type[LFUCache[T, U]], size: int = 128
+ ) -> Callable[[Callable[[T], U]], Callable[..., U]]:
  """
  Decorator version of LFU Cache
+
+ Decorated function must be function of T -> U
  """
 
- def cache_decorator_inner(func: Callable):
- def cache_decorator_wrapper(*args, **kwargs):
- if func not in LFUCache.decorator_function_to_instance_map:
- LFUCache.decorator_function_to_instance_map[func] = LFUCache(size)
+ def cache_decorator_inner(func: Callable[[T], U]) -> Callable[..., U]:
+ def cache_decorator_wrapper(*args: T) -> U:
+ if func not in cls.decorator_function_to_instance_map:
+ cls.decorator_function_to_instance_map[func] = LFUCache(size)
 
- result = LFUCache.decorator_function_to_instance_map[func].get(args[0])
+ result = cls.decorator_function_to_instance_map[func].get(args[0])
  if result is None:
- result = func(*args, **kwargs)
- LFUCache.decorator_function_to_instance_map[func].set(
- args[0], result
- )
+ result = func(*args)
+ cls.decorator_function_to_instance_map[func].set(args[0], result)
  return result
 
- def cache_info():
- return LFUCache.decorator_function_to_instance_map[func]
+ def cache_info() -> LFUCache[T, U]:
+ return cls.decorator_function_to_instance_map[func]
 
- cache_decorator_wrapper.cache_info = cache_info
+ setattr(cache_decorator_wrapper, "cache_info", cache_info)
 
  return cache_decorator_wrapper