The following KeyedSet class mirrors the builtin set class as closely as possible, whilst maintaining the general flexibility of a dictionary. The only requirement is that each set item has a distinct 'id' attribute. The usual set operations are implemented, but items can also be referenced via their id.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 | class KeyedSet(dict):
"""
A set class for handling collections of arbitrary
objects that have unique, and hashable 'id' attributes.
Set items are stored as values in a dictionary, with ids
as keys. There is no requirement for set items to be hashable.
The class requires a 1 to 1 mapping between objects
and their ids, and is designed for cases where access to
items via a key lookup is also desirable.
"""
def __init__(self, items=None):
if items is not None:
for item in items:
self[item.id] = item
def add(self, item):
self[item.id] = item
def remove(self, item):
del self[item.id]
def __contains__(self, item):
try:
return self.has_key(item.id)
except AttributeError:
return False
def __iter__(self):
return self.itervalues()
def __repr__(self):
return '%s(%r)' % (self.__class__.__name__, self.keys())
def __cmp__(self, other):
raise TypeError, "can't compare KeyedSets using cmp()"
def issubset(self, other):
self._binary_check(other)
if len(self) > len(other):
return False
else:
for key in self.iterkeys():
if not other.has_key(key):
return False
return True
def issuperset(self, other):
self._binary_check(other)
return other.issubset(self)
__le__ = issubset
__ge__ = issuperset
def __lt__(self, other):
self._binary_check(other)
return len(self) < len(other) and self.issubset(other)
def __gt__(self, other):
self._binary_check(other)
return len(self) > len(other) and self.issuperset(other)
def __eq__(self, other):
if isinstance(other, self.__class__):
return len(self) == len(other) and self.issubset(other)
else:
return False
def __ne__(self, other):
if isinstance(other, self.__class__):
return not self == other
else:
return True
def union(self, other):
res = self.copy()
for item in other:
res.add(item)
return res
def intersection(self, other):
res = self.__class__()
if not isinstance(other, self.__class__):
other = self.__class__(other)
if len(self) > len(other):
for item in other:
if item in self:
res.add(item)
else:
for item in self:
if item in other:
res.add(item)
return res
def difference(self, other):
res = self.copy()
for item in other:
if item in res:
res.remove(item)
return res
def symmetric_difference(self, other):
res = self.copy()
if not isinstance(other, self.__class__):
other = self.__class__(other)
for item in other:
if item in self:
res.remove(item)
else:
res.add(item)
return res
def __or__(self, other):
self._binary_check(other)
return self.union(other)
def __and__(self, other):
self._binary_check(other)
return self.intersection(other)
def __sub__(self, other):
self._binary_check(other)
return self.difference(other)
def __xor__(self, other):
self._binary_check(other)
return self.symmetric_difference(other)
def _binary_check(self, other):
if not isinstance(other, self.__class__):
raise TypeError, "Binary operation only permitted between KeyedSets"
def copy(self):
res = self.__class__()
res.update(self)
return res
def union_update(self, other):
if isinstance(other, (self.__class__, dict)):
self.update(other)
else:
for item in other:
self.add(item)
def intersection_update(self, other):
if not isinstance(other, self.__class__):
other = self.__class__(other)
self &= other
def difference_update(self, other):
for item in other:
self.discard(item)
def symmetric_difference_update(self, other):
if not isinstance(other, self.__class__):
other = self.__class__(other)
for item in other:
if item in self:
self.remove(item)
else:
self.add(item)
def __ior__(self, other):
self._binary_check(other)
self.union_update(other)
return self
def __iand__(self, other):
self._binary_check(other)
intersect = self & other
self.clear()
self.update(intersect)
return self
def __isub__(self, other):
self._binary_check(other)
self.difference_update(other)
return self
def __ixor__(self, other):
self._binary_check(other)
self.symmetric_difference_update(other)
return self
def discard(self, item):
try:
self.remove(item)
except KeyError:
pass
def pop(self, *args):
if args:
return super(self.__class__, self).pop(*args)
else:
return self.popitem()[1]
def update(self, other):
if isinstance(other, (self.__class__, dict)):
super(self.__class__, self).update(other)
else:
for item in other:
self.add(item)
|
For some applications it is fairly natural to assign distinct objects distinct 'id' attributes. E.g. nodes in a graph. In some cases it is useful (or at least I find it so) to be able to contain such items in sets (and perform set operations) whilst also allowing access via object id.
Use is almost identical to that of the built in set class. The only differences (that I am aware of) are that the pop() method, when supplied with appropriate arguments, calls the dictionary pop() method, rather than popping an arbitrary item; and the __repr__() method returns string representations of the keys (ids), rather than values (items). Almost all dictionary methods are also available. 'k in a' does not work identically to dictionaries as this checks whether item k is in a.values() (or rather that k.id is in a.keys()). Iterating over the set iterates over the values, rather than keys. But the availability of has_key() and iterkeys() ensures that a KeyedSet maintains all the basic functionality of a dictionary when required.