src/dctv/iddict.py

# Copyright (C) 2020 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http:#www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Dictionary types that work for object identity, not defined equality"""
import weakref
from contextlib import contextmanager
from collections.abc import MutableMapping, MutableSet
from threading import RLock

_NO_RELOAD = True


class IdentityDictionary(MutableMapping):
  """Identity version of a dict"""
  def __init__(self):
    self.__data = {}

  def __repr__(self):
    return "ID{%s}" % " ,".join(
      "{!r}: {!r}".format(*item) for item in self.items())

  # Container

  def __contains__(self, key):
    return id(key) in self.__data

  # Mapping

  def __getitem__(self, key):
    return self.__data[id(key)][1]

  def __iter__(self):
    return (v[0] for v in self.__data.values())

  def __len__(self):
    return len(self.__data)

  def keys(self):
    return iter(self)

  def values(self):
    return (v[1] for v in self.items())

  def items(self):
    return self.__data.values()

  # MutableMapping

  def __setitem__(self, key, value):
    self.__data[id(key)] = (key, value)

  def __delitem__(self, key):
    del self.__data[id(key)]

  def clear(self):
    self.__data.clear()


class IdentitySet(MutableSet):
  """Identity version of a set"""
  def __init__(self):
    self.__data = {}

  # Container

  def __contains__(self, item):
    return id(item) in self.__data

  # Set

  def __iter__(self):
    return iter(self.__data.values())

  def __len__(self):
    return len(self.__data)

  # MutableSet

  def add(self, value):
    self.__data[id(value)] = value

  def discard(self, value):
    self.__data.pop(id(value), None)


class WeakKeyIdentityDictionary(MutableMapping):
  """Weakly-keyed dictionary that stores objects by identity"""

  # This class is complicated by how weakref defers to object __hash__
  # instead of object identity.  It works, essentially, by
  # implementing a "hash table" at the Python level, using the object
  # ID as a hash code.  We can collide: we can store an object A with
  # id I that later dies.  In the meantime, we can construct an object
  # B that reuses id I.  So we have two weak refs keyed with id I and
  # have to disambiguate them.  The set of weakrefs to objects with id
  # I is the bucket for I.  The cookie is a unique integer that
  # associates an object with a value.  Each object gets its own
  # cookie, so two objects with the same ID can have different
  # associated data.

  def __init__(self):
    self.__clean_lock = RLock()
    self.__clean_block_depth = 0
    self.__to_clean = set()
    self.__buckets_by_id = {}
    self.__values_by_cookie = {}
    self.__next_cookie = 1  # True in bool context

  @contextmanager
  def __clean_blocked(self):
    with self.__clean_lock:
      self.__clean_block_depth += 1
    yield
    with self.__clean_lock:
      self.__clean_block_depth -= 1
      if not self.__clean_block_depth:
        self.__flush_clean_locked()

  def __flush_clean_locked(self):
    while self.__to_clean:
      bucket_id, cookie = self.__to_clean.pop()
      bucket = self.__buckets_by_id[bucket_id]
      del self.__values_by_cookie[cookie]
      del bucket[cookie]
      if not bucket:
        del self.__buckets_by_id[bucket_id]

  def __weakref_callback(self, wref):
    with self.__clean_lock:
      self.__to_clean.add(wref.key)
      if not self.__clean_block_depth:
        self.__flush_clean_locked()

  def __find(self, key):
    assert key is not None
    assert self.__clean_block_depth > 0  # pylint: disable=compare-to-zero
    bucket = self.__buckets_by_id.get(id(key))
    cookie = None
    for b_cookie, b_ref in bucket.items() if bucket else ():
      candidate = b_ref()
      if candidate is None:
        assert b_ref.key == (id(key), b_cookie)
        self.__to_clean.add(b_ref.key)
      elif candidate is key:
        cookie = b_cookie
        break
    return cookie, bucket

  # Container

  def __contains__(self, key):
    with self.__clean_blocked():
      cookie, _ = self.__find(key)
      return bool(cookie)

  # Mapping

  def __getitem__(self, key):
    with self.__clean_blocked():
      cookie, _ = self.__find(key)
      if cookie:
        return self.__values_by_cookie[cookie]
      raise KeyError(key)

  def get(self, key, default=None):
    with self.__clean_blocked():
      cookie, _ = self.__find(key)
      if cookie:
        return self.__values_by_cookie[cookie]
      return default

  def __iter__(self):
    return iter(self.keys())

  def __len__(self):
    return len(self.items())

  def keys(self):
    return (v[0] for v in self.items())

  def values(self):
    return (v[1] for v in self.items())

  def items(self):
    # Weak-keyed collections are inherently unstable, so just return a
    # snapshot of the object contents.
    with self.__clean_blocked():
      items = []
      for bucket in self.__buckets_by_id.values():
        for b_cookie, b_ref in bucket.items():
          obj = b_ref()
          if obj is None:
            self.__to_clean.add(b_ref.key)
          else:
            item = (obj, self.__values_by_cookie[b_cookie])
            items.append(item)
      return items

  # MutableMapping

  def __setitem__(self, key, value):
    with self.__clean_blocked():
      cookie, bucket = self.__find(key)
      if not cookie:
        if not bucket:
          bucket = {}
          self.__buckets_by_id[id(key)] = bucket
        cookie = self.__next_cookie
        bucket[cookie] = weakref.KeyedRef(
          key,
          self.__weakref_callback,
          (id(key), cookie))
        self.__next_cookie += 1
      self.__values_by_cookie[cookie] = value

  def __delitem__(self, key):
    with self.__clean_blocked():
      cookie, _ = self.__find(key)
      if not cookie:
        raise KeyError(key)
      self.__to_clean.add((id(key), cookie))