← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/intset.py

  • Committer: Martin Pool
  • Date: 2005-04-28 07:24:55 UTC
  • Revision ID: mbp@sourcefrog.net-20050428072453-7b99afa993a1e549
todo

Show diffs side-by-side

added added

removed removed

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