~bzr-pqm/bzr/bzr.dev

1185.1.29 by Robert Collins
merge merge tweaks from aaron, which includes latest .dev
1
# Copyright (C) 2005 Canonical Ltd
1887.1.1 by Adeodato Simó
Do not separate paragraphs in the copyright statement with blank lines,
2
#
1185.1.29 by Robert Collins
merge merge tweaks from aaron, which includes latest .dev
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.
1887.1.1 by Adeodato Simó
Do not separate paragraphs in the copyright statement with blank lines,
7
#
1185.1.29 by Robert Collins
merge merge tweaks from aaron, which includes latest .dev
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.
1887.1.1 by Adeodato Simó
Do not separate paragraphs in the copyright statement with blank lines,
12
#
1185.1.29 by Robert Collins
merge merge tweaks from aaron, which includes latest .dev
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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
16
17
# Author: Martin Pool <mbp@canonical.com>
18
19
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.
23
24
25
class IntSet(Exception):
26
    """Faster set-like class storing only whole numbers.
27
28
    Despite the name this stores long integers happily, but negative
29
    values are not allowed.
30
31
    >>> a = IntSet([0, 2, 5])
32
    >>> bool(a)
33
    True
34
    >>> 2 in a
35
    True
36
    >>> 4 in a
37
    False
38
    >>> a.add(4)
39
    >>> 4 in a
40
    True
41
42
    >>> b = IntSet()
43
    >>> not b
44
    True
45
    >>> b.add(10)
46
    >>> 10 in a
47
    False
48
    >>> a.update(b)
49
    >>> 10 in a
50
    True
51
    >>> a.update(range(5))
52
    >>> 3 in a
53
    True
54
55
    Being a set, duplicates are ignored:
56
    >>> a = IntSet()
57
    >>> a.add(10)
58
    >>> a.add(10)
59
    >>> 10 in a
60
    True
61
    >>> list(a)
62
    [10]
63
    
64
    """
65
    __slots__ = ['_val']
66
67
    def __init__(self, values=None, bitmask=0L):
68
        """Create a new intset.
69
70
        values
71
            If specified, an initial collection of values.
72
        """
73
        self._val = bitmask
1963.2.6 by Robey Pointer
pychecker is on crack; go back to using 'is None'.
74
        if values is not None:
1185.1.29 by Robert Collins
merge merge tweaks from aaron, which includes latest .dev
75
            self.update(values)
76
77
78
    def __nonzero__(self):
79
        """IntSets are false if empty, otherwise True.
80
81
        >>> bool(IntSet())
82
        False
83
        
84
        >>> bool(IntSet([0]))
85
        True
86
        """
87
        return bool(self._val)
88
89
90
    def __len__(self):
91
        """Number of elements in set.
92
93
        >>> len(IntSet(xrange(20000)))
94
        20000
95
        """
96
        v = self._val
97
        c = 0
98
        while v:
99
            if v & 1:
100
                c += 1
101
            v = v >> 1
102
        return c
103
104
105
    def __and__(self, other):
106
        """Set intersection.
107
108
        >>> a = IntSet(range(10))
109
        >>> len(a)
110
        10
111
        >>> b = a & a
112
        >>> b == a
113
        True
114
        >>> a = a & IntSet([5, 7, 11, 13])
115
        >>> list(a)
116
        [5, 7]
117
        """
118
        if not isinstance(other, IntSet):
119
            raise NotImplementedError(type(other))
120
        return IntSet(bitmask=(self._val & other._val))
121
122
123
    def __or__(self, other):
124
        """Set union.
125
126
        >>> a = IntSet(range(10)) | IntSet([5, 15, 25])
127
        >>> len(a)
128
        12
129
        """
130
        if not isinstance(other, IntSet):
131
            raise NotImplementedError(type(other))
132
        return IntSet(bitmask=(self._val | other._val))        
133
134
135
    def __eq__(self, other):
136
        """Comparison.
137
138
        >>> IntSet(range(3)) == IntSet([2, 0, 1])
139
        True
140
        """
141
        if isinstance(other, IntSet):
142
            return self._val == other._val
143
        else:
144
            return False
145
146
147
    def __ne__(self, other):
148
        return not self.__eq__(other)
149
150
151
    def __contains__(self, i):
152
        assert i >= 0
153
        return self._val & (1L << i)
154
155
156
    def __iter__(self):
157
        """Return contents of set.
158
159
        >>> list(IntSet())
160
        []
161
        >>> list(IntSet([0, 1, 5, 7]))
162
        [0, 1, 5, 7]
163
        """
164
        v = self._val
165
        o = 0
166
        # XXX: This is a bit slow
167
        while v:
168
            if v & 1:
169
                yield o
170
            v = v >> 1
171
            o = o + 1
172
173
        
174
    def update(self, to_add):
175
        """Add all the values from the sequence or intset to_add"""
176
        if isinstance(to_add, IntSet):
177
            self._val |= to_add._val
178
        else:
179
            for i in to_add:
180
                assert i >= 0
181
                self._val |= (1L << i)
182
183
184
    def add(self, to_add):
185
        assert 0 <= to_add
186
        self._val |= (1L << to_add)
187
188
189
    def remove(self, to_remove):
190
        """Remove one value from the set.
191
192
        Raises KeyError if the value is not present.
193
194
        >>> a = IntSet([10])
195
        >>> a.remove(9)
196
        Traceback (most recent call last):
197
          File "/usr/lib/python2.4/doctest.py", line 1243, in __run
198
            compileflags, 1) in test.globs
199
          File "<doctest __main__.IntSet.remove[1]>", line 1, in ?
200
            a.remove(9)
201
        KeyError: 9
202
        >>> a.remove(10)
203
        >>> not a
204
        True
205
        """
206
        assert 0 <= to_remove
207
        m = 1L << to_remove
208
        if not self._val & m:
209
            raise KeyError(to_remove)
210
        self._val ^= m
1185.50.24 by John Arbash Meinel
Added set_remove to get the set delete function to IntSet
211
212
    def set_remove(self, to_remove):
213
        """Remove all values that exist in to_remove.
214
215
        >>> a = IntSet(range(10))
216
        >>> b = IntSet([2,3,4,7,12])
217
        >>> a.set_remove(b)
218
        >>> list(a)
219
        [0, 1, 5, 6, 8, 9]
220
        >>> a.set_remove([1,2,5])
221
        >>> list(a)
222
        [0, 6, 8, 9]
223
        """
224
        if not isinstance(to_remove, IntSet):
225
            self.set_remove(IntSet(to_remove))
226
            return
227
        intersect = self._val & to_remove._val
228
        self._val ^= intersect
229