from google.appengine.ext import db

class EntityCache(object):
  """Caches fetched entities so subsequent fetches return the same instance.

  Usage:
  cache = EntityCache()
  e1, e2 = cache.get([k1, k2])
  e3, e4 = cache.get([k2, k3])
  assert e2 == e3
  """

  def __init__(self):
    self.cache = {}

  def get(self, keys):
    """Get the entities matching the provided keys."""
    if isinstance(keys, db.Key):
      keys = [keys]
    missing_keys = [k for k in keys if k not in self.cache]
    if missing_keys:
      self.cache.update((x.key(), x) for x in db.get(missing_keys))
    return [self.cache[k] for k in keys]