1
 
# Copyright (C) 2005 Canonical Ltd
 
3
 
# This program is free software; you can redistribute it and/or modify
 
4
 
# it under the terms of the GNU General Public License as published by
 
5
 
# the Free Software Foundation; either version 2 of the License, or
 
6
 
# (at your option) any later version.
 
8
 
# This program is distributed in the hope that it will be useful,
 
9
 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
10
 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
11
 
# GNU General Public License for more details.
 
13
 
# You should have received a copy of the GNU General Public License
 
14
 
# along with this program; if not, write to the Free Software
 
15
 
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 
17
 
# Author: Martin Pool <mbp@canonical.com>
 
20
 
# Somewhat surprisingly, it turns out that this is much slower than
 
21
 
# simply storing the ints in a set() type.  Python's performance model
 
22
 
# is very different to that of C.
 
25
 
class IntSet(Exception):
 
26
 
    """Faster set-like class storing only whole numbers.
 
28
 
    Despite the name this stores long integers happily, but negative
 
29
 
    values are not allowed.
 
31
 
    >>> a = IntSet([0, 2, 5])
 
51
 
    >>> a.update(range(5))
 
55
 
    Being a set, duplicates are ignored:
 
67
 
    def __init__(self, values=None, bitmask=0L):
 
68
 
        """Create a new intset.
 
71
 
            If specified, an initial collection of values.
 
74
 
        if values is not None:
 
78
 
    def __nonzero__(self):
 
79
 
        """IntSets are false if empty, otherwise True.
 
87
 
        return bool(self._val)
 
91
 
        """Number of elements in set.
 
93
 
        >>> len(IntSet(xrange(20000)))
 
105
 
    def __and__(self, other):
 
108
 
        >>> a = IntSet(range(10))
 
114
 
        >>> a = a & IntSet([5, 7, 11, 13])
 
118
 
        if not isinstance(other, IntSet):
 
119
 
            raise NotImplementedError(type(other))
 
120
 
        return IntSet(bitmask=(self._val & other._val))
 
123
 
    def __or__(self, other):
 
126
 
        >>> a = IntSet(range(10)) | IntSet([5, 15, 25])
 
130
 
        if not isinstance(other, IntSet):
 
131
 
            raise NotImplementedError(type(other))
 
132
 
        return IntSet(bitmask=(self._val | other._val))
 
135
 
    def __eq__(self, other):
 
138
 
        >>> IntSet(range(3)) == IntSet([2, 0, 1])
 
141
 
        if isinstance(other, IntSet):
 
142
 
            return self._val == other._val
 
147
 
    def __ne__(self, other):
 
148
 
        return not self.__eq__(other)
 
151
 
    def __contains__(self, i):
 
152
 
        return self._val & (1L << i)
 
156
 
        """Return contents of set.
 
160
 
        >>> list(IntSet([0, 1, 5, 7]))
 
165
 
        # XXX: This is a bit slow
 
173
 
    def update(self, to_add):
 
174
 
        """Add all the values from the sequence or intset to_add"""
 
175
 
        if isinstance(to_add, IntSet):
 
176
 
            self._val |= to_add._val
 
179
 
                self._val |= (1L << i)
 
182
 
    def add(self, to_add):
 
183
 
        self._val |= (1L << to_add)
 
186
 
    def remove(self, to_remove):
 
187
 
        """Remove one value from the set.
 
189
 
        Raises KeyError if the value is not present.
 
193
 
        Traceback (most recent call last):
 
194
 
          File "/usr/lib/python2.4/doctest.py", line 1243, in __run
 
195
 
            compileflags, 1) in test.globs
 
196
 
          File "<doctest __main__.IntSet.remove[1]>", line 1, in ?
 
204
 
        if not self._val & m:
 
205
 
            raise KeyError(to_remove)
 
208
 
    def set_remove(self, to_remove):
 
209
 
        """Remove all values that exist in to_remove.
 
211
 
        >>> a = IntSet(range(10))
 
212
 
        >>> b = IntSet([2,3,4,7,12])
 
216
 
        >>> a.set_remove([1,2,5])
 
220
 
        if not isinstance(to_remove, IntSet):
 
221
 
            self.set_remove(IntSet(to_remove))
 
223
 
        intersect = self._val & to_remove._val
 
224
 
        self._val ^= intersect