Search code examples
pythoninheritanceweak-references

Indexable weak ordered set in Python


I was wondering if there is an easy way to build an indexable weak ordered set in Python. I tried to build one myself. Here's what I came up with:

"""
An indexable, ordered set of objects, which are held by weak reference.
"""
from nose.tools import *
import blist
import weakref


class WeakOrderedSet(blist.weaksortedset):
    """
    A blist.weaksortedset whose key is the insertion order.
    """
    def __init__(self, iterable=()):
        self.insertion_order = weakref.WeakKeyDictionary()  # value_type to int
        self.last_key = 0
        super().__init__(key=self.insertion_order.__getitem__)
        for item in iterable:
            self.add(item)

    def __delitem__(self, index):
        values = super().__getitem__(index)
        super().__delitem__(index)
        if not isinstance(index, slice):
            # values is just one element
            values = [values]
        for value in values:
            if value not in self:
                del self.insertion_order[value]

    def add(self, value):
        # Choose a key so that value is on the end.
        if value not in self.insertion_order:
            key = self.last_key
            self.last_key += 1
            self.insertion_order[value] = key
        super().add(value)

    def discard(self, value):
        super().discard(value)
        if value not in self:
            del self.insertion_order[value]

    def remove(self, value):
        super().remove(value)
        if value not in self:
            del self.insertion_order[value]

    def pop(self, *args, **kwargs):
        value = super().pop(*args, **kwargs)
        if value not in self:
            del self.insertion_order[value]

    def clear(self):
        super().clear()
        self.insertion_order.clear()

    def update(self, *args):
        for arg in args:
            for item in arg:
                self.add(item)


if __name__ == '__main__':
    class Dummy:
        def __init__(self, value):
            self.value = value

    x = [Dummy(i) for i in range(10)]
    w = WeakOrderedSet(reversed(x))
    del w[2:8]
    assert_equals([9,8,1,0], [i.value for i in w])
    del w[0]
    assert_equals([8,1,0], [i.value for i in w])
    del x
    assert_equals([], [i.value for i in w])

Is there an easier way to do this?


Solution

  • The easiest way to is to take advantage of existing components in the standard library.

    OrderedDict and the MutableSet ABC make it easy to write an OrderedSet.

    Likewise, you can reuse the existing weakref.WeakSet and replace its underlying set() with an OrderedSet.

    Indexing is more difficult to achieve -- these easiest way it to convert it to a list when needed. That is necessary because sets and dicts are intrinsically sparse.

    import collections.abc
    import weakref
    
    class OrderedSet(collections.abc.MutableSet):
        def __init__(self, values=()):
            self._od = collections.OrderedDict().fromkeys(values)
        def __len__(self):
            return len(self._od)
        def __iter__(self):
            return iter(self._od)
        def __contains__(self, value):
            return value in self._od
        def add(self, value):
            self._od[value] = None
        def discard(self, value):
            self._od.pop(value, None)
    
    class OrderedWeakrefSet(weakref.WeakSet):
        def __init__(self, values=()):
            super(OrderedWeakrefSet, self).__init__()
            self.data = OrderedSet()
            for elem in values:
                self.add(elem)
    

    Use it like this:

    >>> names = OrderedSet(['Alice', 'Bob', 'Carol', 'Bob', 'Dave', 'Edna'])
    >>> len(names)
    5
    >>> 'Bob' in names
    True
    >>> s = list(names)
    >>> s[2]
    'Carol'
    >>> s[4]
    'Edna'
    

    Note as of Python 3.7, regular dicts are guaranteed to be ordered, so you can substitute dict for OrderedDict in this recipe and it will all work fine :-)