lru.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import threading


class Node(object):
    def __init__(self,
                 key, value,
                 pre=None, next=None):
        self.key = key
        self.value = value
        self.pre = pre
        self.next = next


class LRUCache(object):
    def __init__(self, capacity):
        if type(capacity) != int and capacity <= 0:
            raise Exception("lru cache max size must be int and larger than 0,"
                            " now is: %s" % capacity)
        self.capacity = capacity

        # init double way linked list
        self._init_list()

        # init length
        self.length = 0

        # init data set
        self._init_dataset()

        # init lock
        self.lock = threading.Lock()

    def _init_list(self):
        self.head = Node(key=None,
                         value=None)
        self.tail = Node(key=None,
                         value=None)
        self.head.next = self.tail
        self.tail.pre = self.head

    def _init_dataset(self):
        self.hash_map = dict()

    def _get_node(self, key):
        return self.hash_map.get(key, None)

    def get(self, key):
        with self.lock:
            node = self._get_node(key)
            if node is not None:
                self._mv_to_first(node)
                return node.value
            return None

    def _mv_to_first(self, node):
        node.next.pre = node.pre
        node.pre.next = node.next
        node.next = self.head.next
        self.head.next.pre = node
        self.head.next = node
        node.pre = self.head

    def set(self, key, value):
        with self.lock:
            if self.has_key(key):
                node = self._get_node(key)
                node.value = value
                self._mv_to_first(node)
            else:
                # if it is full, remove the last one.
                if self.full():
                    d_node = self.tail.pre
                    self.tail.pre = d_node.pre
                    d_node.pre.next = self.tail
                    d_node.next = None
                    d_node.pre = None
                    del d_node
                    self.length -= 1
                node = Node(key, value, self.head, self.head.next)
                self.head.next.pre = node
                self.head.next = node
                self.length += 1
            self.hash_map[key] = node

    def traverse(self):
        node = self.head.next
        while node != self.tail:
            yield (node.key, node.value)
            node = node.next

    def full(self):
        return self.length >= self.capacity

    def has_key(self, key):
        node = self._get_node(key)
        return not (node is None)


if __name__ == '__main__':
    cache = LRUCache(5)
    cache.set(1, 1)
    cache.set(2, 2)
    cache.set(3, 3)
    cache.set(4, 4)
    cache.set(5, 5)
    for key, value in cache.traverse():
        print key, value
    print

    cache.get(1)
    for key, value in cache.traverse():
        print key, value
    print

    cache.set(6, 6)
    for key, value in cache.traverse():
        print key, value
    print

    cache.set(3, 3)
    for key, value in cache.traverse():
        print key, value
    print

    print cache.has_key(4), cache.get(4)
    print cache.has_key(7), cache.get(7)