quay/util/expiresdict.py

from datetime import datetime

class ExpiresEntry(object):
  """ A single entry under a ExpiresDict. """
  def __init__(self, value, expires=None):
    self.value = value
    self._expiration = expires

  @property
  def expired(self):
    if self._expiration is None:
      return False

    return datetime.now() >= self._expiration


class ExpiresDict(object):
  """ ExpiresDict defines a dictionary-like class whose keys have expiration. The rebuilder is
      a function that returns the full contents of the cached dictionary as a dict of the keys
      and whose values are TTLEntry's. If the rebuilder is None, then no rebuilding is performed.
  """
  def __init__(self, rebuilder=None):
    self._rebuilder = rebuilder
    self._items = {}

  def __getitem__(self, key):
    found = self.get(key)
    if found is None:
      raise KeyError

    return found

  def get(self, key, default_value=None):
    # Check the cache first. If the key is found and it has not yet expired,
    # return it.
    found = self._items.get(key)
    if found is not None and not found.expired:
      return found.value

    # Otherwise the key has expired or was not found. Rebuild the cache and check it again.
    items = self._rebuild()
    found_item = items.get(key)
    if found_item is None:
      return default_value

    return found_item.value

  def __contains__(self, key):
    return self.get(key) is not None

  def _rebuild(self):
    if self._rebuilder is None:
      return self._items

    items = self._rebuilder()
    self._items = items
    return items

  def set(self, key, value, expires=None):
    self._items[key] = ExpiresEntry(value, expires=expires)