[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

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.freq: int = 0
-        self.next = None
+        self.next: DoubleLinkedListNode[T, U] | None = None
-        self.prev = 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):
+    def __init__(self) -> None:
-        self.head = DoubleLinkedListNode(None, None)
+        self.head: DoubleLinkedListNode[T, U] = DoubleLinkedListNode(None, None)
-        self.rear = 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
 
-        self.rear.prev, node.next = node, self.rear
+        # All nodes other than self.head are guaranteed to have non-None previous
-        temp.next, node.prev = node, temp
+        assert previous is not None
+
+        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:
+    def _position_node(self, node: DoubleLinkedListNode[T, U]) -> None:
-        while node.prev.key and node.prev.freq > node.freq:
+        """
-            node1, node2 = node, node.prev
+        Moves node forward to maintain invariant of sort by freq value
-            node1.prev, node2.next = node2.prev, node1.prev
+        """
-            node1.next, node2.prev = node2, node1
 
-    def remove(self, node: DoubleLinkedListNode) -> DoubleLinkedListNode:
+        while node.prev is not None and node.prev.freq > node.freq:
+            # swap node with previous node
+            previous_node = node.prev
+
+            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
+        if node.prev is None or node.next is None:
-        node.prev, node.next = None, None
+            return None
-        temp_last.next, temp_next.prev = temp_next, temp_last
+
+        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,40 +216,57 @@ class LFUCache:
             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]))
+            value_node: DoubleLinkedListNode[T, U] = self.cache[key]
-            return self.cache[key].val
+            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
+                # delete first node when over capacity
-                self.list.remove(self.cache[key_to_delete])
+                first_node = self.list.head.next
-                del self.cache[key_to_delete]
+
+                # 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])
@@ -151,32 +274,35 @@ class LFUCache:
 
         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
+    @classmethod
-    def decorator(size: int = 128):
+    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_inner(func: Callable[[T], U]) -> Callable[..., U]:
-            def cache_decorator_wrapper(*args, **kwargs):
+            def cache_decorator_wrapper(*args: T) -> U:
-                if func not in LFUCache.decorator_function_to_instance_map:
+                if func not in cls.decorator_function_to_instance_map:
-                    LFUCache.decorator_function_to_instance_map[func] = LFUCache(size)
+                    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)
+                    result = func(*args)
-                    LFUCache.decorator_function_to_instance_map[func].set(
+                    cls.decorator_function_to_instance_map[func].set(args[0], result)
-                        args[0], result
-                    )
                 return result
 
-            def cache_info():
+            def cache_info() -> LFUCache[T, U]:
-                return LFUCache.decorator_function_to_instance_map[func]
+                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
 

other/lru_cache.py

@@ -1,52 +1,147 @@
 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 LRU Cache
+
+    >>> DoubleLinkedListNode(1,1)
+    Node: key: 1, val: 1, 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.next = None
+        self.next: DoubleLinkedListNode[T, U] | None = None
-        self.prev = None
+        self.prev: DoubleLinkedListNode[T, U] | None = None
+
+    def __repr__(self) -> str:
+        return "Node: key: {}, val: {}, has next: {}, has prev: {}".format(
+            self.key, self.val, self.next is not None, self.prev is not None
+        )
 
 
-class DoubleLinkedList:
+class DoubleLinkedList(Generic[T, U]):
     """
     Double Linked List built specifically for LRU Cache
+
+    >>> dll: DoubleLinkedList = DoubleLinkedList()
+    >>> dll
+    DoubleLinkedList,
+        Node: key: None, val: None, has next: True, has prev: False,
+        Node: key: None, val: None, has next: False, has prev: True
+
+    >>> first_node = DoubleLinkedListNode(1,10)
+    >>> first_node
+    Node: key: 1, val: 10, has next: False, has prev: False
+
+
+    >>> dll.add(first_node)
+    >>> dll
+    DoubleLinkedList,
+        Node: key: None, val: None, has next: True, has prev: False,
+        Node: key: 1, val: 10, has next: True, has prev: True,
+        Node: key: None, val: None, has next: False, has prev: True
+
+    >>> # node is mutated
+    >>> first_node
+    Node: key: 1, val: 10, has next: True, has prev: True
+
+    >>> second_node = DoubleLinkedListNode(2,20)
+    >>> second_node
+    Node: key: 2, val: 20, has next: False, has prev: False
+
+    >>> dll.add(second_node)
+    >>> dll
+    DoubleLinkedList,
+        Node: key: None, val: None, has next: True, has prev: False,
+        Node: key: 1, val: 10, has next: True, has prev: True,
+        Node: key: 2, val: 20, has next: True, has prev: True,
+        Node: key: None, val: None, has next: False, has prev: True
+
+    >>> removed_node = dll.remove(first_node)
+    >>> assert removed_node == first_node
+    >>> dll
+    DoubleLinkedList,
+        Node: key: None, val: None, has next: True, has prev: False,
+        Node: key: 2, val: 20, has next: True, has prev: True,
+        Node: key: None, val: None, 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, 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, has next: False, has prev: True
+    >>> dll.remove(dll.rear) is None
+    True
+
+
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
-        self.head = DoubleLinkedListNode(None, None)
+        self.head: DoubleLinkedListNode[T, U] = DoubleLinkedListNode(None, None)
-        self.rear = 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 to the end of the list (before rear)
         """
 
-        temp = self.rear.prev
+        previous = self.rear.prev
-        temp.next, node.prev = node, temp
-        self.rear.prev, node.next = node, self.rear
 
-    def remove(self, node: DoubleLinkedListNode) -> DoubleLinkedListNode:
+        # All nodes other than self.head are guaranteed to have non-None previous
+        assert previous is not None
+
+        previous.next = node
+        node.prev = previous
+        self.rear.prev = node
+        node.next = self.rear
+
+    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
+        if node.prev is None or node.next is None:
-        node.prev, node.next = None, None
+            return None
-        temp_last.next, temp_next.prev = temp_next, temp_last
 
+        node.prev.next = node.next
+        node.next.prev = node.prev
+        node.prev = None
+        node.next = None
         return node
 
 
-class LRUCache:
+class LRUCache(Generic[T, U]):
     """
     LRU Cache to store a given capacity of data. Can be used as a stand-alone object
     or as a function decorator.
@@ -54,19 +149,41 @@ class LRUCache:
     >>> cache = LRUCache(2)
 
     >>> cache.set(1, 1)
-
     >>> cache.set(2, 2)
-
     >>> cache.get(1)
     1
 
+    >>> cache.list
+    DoubleLinkedList,
+        Node: key: None, val: None, has next: True, has prev: False,
+        Node: key: 2, val: 2, has next: True, has prev: True,
+        Node: key: 1, val: 1, has next: True, has prev: True,
+        Node: key: None, val: None, has next: False, has prev: True
+
+    >>> cache.cache  # doctest: +NORMALIZE_WHITESPACE
+    {1: Node: key: 1, val: 1, has next: True, has prev: True, \
+     2: Node: key: 2, val: 2, has next: True, has prev: True}
+
     >>> cache.set(3, 3)
 
-    >>> cache.get(2)    # None returned
+    >>> cache.list
+    DoubleLinkedList,
+        Node: key: None, val: None, has next: True, has prev: False,
+        Node: key: 1, val: 1, has next: True, has prev: True,
+        Node: key: 3, val: 3, has next: True, has prev: True,
+        Node: key: None, val: None, has next: False, has prev: True
+
+    >>> cache.cache  # doctest: +NORMALIZE_WHITESPACE
+    {1: Node: key: 1, val: 1, has next: True, has prev: True, \
+     3: Node: key: 3, val: 3, has next: True, has prev: True}
+
+    >>> cache.get(2) is None
+    True
 
     >>> cache.set(4, 4)
 
-    >>> cache.get(1)    # None returned
+    >>> cache.get(1) is None
+    True
 
     >>> cache.get(3)
     3
@@ -91,15 +208,15 @@ class LRUCache:
     """
 
     # 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], LRUCache[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:
         """
@@ -112,7 +229,7 @@ class LRUCache:
             f"capacity={self.capacity}, current size={self.num_keys})"
         )
 
-    def __contains__(self, key: int) -> bool:
+    def __contains__(self, key: T) -> bool:
         """
         >>> cache = LRUCache(1)
 
@@ -127,62 +244,82 @@ class LRUCache:
 
         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
+        Returns the value for the input key and updates the Double Linked List.
-        None if key is not present in cache
+        Returns None if key is not present in cache
         """
+        # Note: pythonic interface would throw KeyError rather than return None
 
         if key in self.cache:
             self.hits += 1
-            self.list.add(self.list.remove(self.cache[key]))
+            value_node: DoubleLinkedListNode[T, U] = self.cache[key]
-            return self.cache[key].val
+            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
+                # delete first node (oldest) when over capacity
-                self.list.remove(self.cache[key_to_delete])
+                first_node = self.list.head.next
-                del self.cache[key_to_delete]
+
+                # 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
+                )  # node guaranteed to be in list assert node.key is not None
+
+                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:
+            # bump node to the end of the list, update value
             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
+    @classmethod
-    def decorator(size: int = 128):
+    def decorator(
+        cls, size: int = 128
+    ) -> Callable[[Callable[[T], U]], Callable[..., U]]:
         """
         Decorator version of LRU Cache
+
+        Decorated function must be function of T -> U
         """
 
-        def cache_decorator_inner(func: Callable):
+        def cache_decorator_inner(func: Callable[[T], U]) -> Callable[..., U]:
-            def cache_decorator_wrapper(*args, **kwargs):
+            def cache_decorator_wrapper(*args: T) -> U:
-                if func not in LRUCache.decorator_function_to_instance_map:
+                if func not in cls.decorator_function_to_instance_map:
-                    LRUCache.decorator_function_to_instance_map[func] = LRUCache(size)
+                    cls.decorator_function_to_instance_map[func] = LRUCache(size)
 
-                result = LRUCache.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)
+                    result = func(*args)
-                    LRUCache.decorator_function_to_instance_map[func].set(
+                    cls.decorator_function_to_instance_map[func].set(args[0], result)
-                        args[0], result
-                    )
                 return result
 
-            def cache_info():
+            def cache_info() -> LRUCache[T, U]:
-                return LRUCache.decorator_function_to_instance_map[func]
+                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