├── .gitignore ├── .travis.yml ├── LICENSE ├── README.rst ├── perftest.py ├── setup.py ├── sortedsets.py └── test_sortedsets.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *.pyo 3 | *.pyc 4 | /build 5 | /dist 6 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | # - "2.7" # unittest.mock doesn't work, may be add fallback? 4 | # - "3.2" # same as 2.7 5 | - "3.3" 6 | - "3.4" 7 | - "3.5" 8 | script: nosetests 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2012 Paul Colomiets 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in 11 | all copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | THE SOFTWARE. 20 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | ========== 2 | SortedSets 3 | ========== 4 | 5 | 6 | Sorted set is data structure that maps unique keys to scores. The following 7 | operations are allowed on sorted set: 8 | 9 | * Iteration in order or scores (both ascending and descending) 10 | * Get score for key (same performance as for dict) 11 | * Get index for key (O(log n), n is sizeof the set) 12 | * Get key for index (O(log n)) 13 | * Slicing by index (O(m + log n), m is length of slice) 14 | * Slicing by score (O(m + log n), m is length of slice) 15 | * Item/slice deletion by index and score (same performance as for slicing) 16 | * Insertion with any score has O(log n) performance too 17 | 18 | The data structure is modelled closely after Redis' sorted sets. Internally it 19 | consists of a mapping between keys and scores, and a skiplist for scores. 20 | 21 | The use cases for SortedSets are following: 22 | 23 | * Leaderboard for a game 24 | * Priority queue (that support task deletion) 25 | * Timer list (supports deletion too) 26 | * Caches with TTL-based, LFU or LRU eviction 27 | * Search databases with relevance scores 28 | * Statistics 29 | * Replacement for ``collections.Counter`` with faster ``most_common()`` 30 | 31 | 32 | Example Code 33 | ============ 34 | 35 | Let's model a leaderboard:: 36 | 37 | >>> from sortedsets import SortedSet 38 | >>> ss = SortedSet() 39 | 40 | Insert some players into sortedset (with some strange made up scores):: 41 | 42 | >>> for i in range(1, 1000): 43 | ... ss['player' + str(i)] = i*10 if i % 2 else i*i 44 | ... 45 | 46 | Let's find out score for player:: 47 | 48 | >>> ss['player20'], ss['player21'] 49 | (400, 210) 50 | 51 | Let's find out their rating positions:: 52 | 53 | >>> ss.index('player20'), ss.index('player21') 54 | (29, 17) 55 | 56 | Let's find out players that have score similar to one's:: 57 | 58 | >>> ss['player49'] 59 | 490 60 | >>> ss.by_score[470:511] 61 | 62 | >>> for k, v in _.items(): 63 | ... print(k, v) 64 | ... 65 | player47 470 66 | player22 484 67 | player49 490 68 | player51 510 69 | 70 | Let's find out players on the rating page 25:: 71 | 72 | >>> page, pagesize = 25, 10 73 | >>> ss.by_index[page*pagesize:page*pagesize + pagesize] 74 | 75 | >>> len(_) 76 | 10 77 | 78 | -------------------------------------------------------------------------------- /perftest.py: -------------------------------------------------------------------------------- 1 | import resource 2 | from time import clock 3 | 4 | from sortedsets import SortedSet 5 | 6 | def test(size): 7 | tm = clock() 8 | ss = SortedSet((str(i), i*10) for i in range(size)) 9 | create_time = clock() - tm 10 | print("SORTED SET WITH", size, "ELEMENTS", ss._level, "LEVELS") 11 | print("Memory usage", resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) 12 | print("Creation time ", format(create_time, '10.2f'), "s") 13 | num = 1000 14 | step = size // (num + 2) 15 | items = [] 16 | for i in range(step, size-step, step): 17 | items.append((str(i), i*10)) 18 | tm = clock() 19 | for k, v in items: 20 | del ss[k] 21 | del_time = num/(clock() - tm) 22 | tm = clock() 23 | for k, v in items: 24 | ss[k] = v 25 | ins_time = num/(clock() - tm) 26 | print("Insertion speed", format(ins_time, '10.2f'), "ins/s") 27 | print("Deletion speed ", format(del_time, '10.2f'), "del/s") 28 | 29 | for size in (10000, 100000, 1000000, 10000000): 30 | test(size) 31 | 32 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | 3 | setup(name='sortedsets', 4 | version='1.0', 5 | description="SortedSet structure closely modelled after Redis' sorted sets", 6 | author='Paul Colomiets', 7 | author_email='paul@colomiets.name', 8 | url='http://github.com/tailhook/sortedsets', 9 | classifiers=[ 10 | 'Programming Language :: Python :: 3', 11 | 'License :: OSI Approved :: MIT License', 12 | ], 13 | py_modules=[ 14 | 'sortedsets', 15 | ], 16 | ) 17 | -------------------------------------------------------------------------------- /sortedsets.py: -------------------------------------------------------------------------------- 1 | import random 2 | import reprlib 3 | from collections import namedtuple, MutableMapping 4 | from itertools import islice 5 | 6 | 7 | empty = object() 8 | 9 | 10 | class Pointer: 11 | __slots__ = ('forward', 'span') 12 | 13 | def __init__(self, forward=None, span=0): 14 | self.forward = forward 15 | self.span = span 16 | 17 | def __repr__(self): 18 | if self.forward: 19 | return '

'.format(self.span, self.forward.key) 20 | else: 21 | # assert self.span == 0 22 | return '

' 23 | 24 | 25 | class Item: 26 | __slots__ = ('key', 'score', 'pointers', 'backward') 27 | 28 | def __init__(self, key, score): 29 | self.key = key 30 | self.score = score 31 | self.pointers = [] 32 | self.backward = None 33 | 34 | def __getitem__(self, i): 35 | if i == len(self.pointers): 36 | self.pointers.append(Pointer()) 37 | elif i > len(self.pointers): 38 | raise RuntimeError("One level at a time required") 39 | return self.pointers[i] 40 | 41 | def __lt__(self, other): 42 | return (self.score < other.score or 43 | self.score == other.score and hash(self.key) < hash(other.key)) 44 | 45 | def __le__(self, other): 46 | return (self.score < other.score or 47 | self.score == other.score and 48 | (hash(self.key) < hash(other.key) or self.key == other.key)) 49 | 50 | def __repr__(self): 51 | return ''.format( 52 | self.key, self.score, self.pointers) 53 | 54 | def _iter_to(self, stop_item): 55 | item = self 56 | while item != stop_item: 57 | assert item is not None 58 | yield item 59 | item = item[0].forward 60 | 61 | def _iter_backwards_to(self, stop_item): 62 | item = self 63 | while item != stop_item: 64 | assert item is not None 65 | yield item 66 | item = item.backward 67 | 68 | 69 | class RankView: 70 | __slots__ = ('_set',) 71 | 72 | def __init__(self, set): 73 | self._set = set 74 | 75 | def __getitem__(self, key): 76 | if isinstance(key, slice): 77 | set_len = len(self._set) 78 | start, stop, step = key.indices(set_len) 79 | 80 | if step <= 0: 81 | raise ValueError("Negative step is useless") 82 | if stop <= start: 83 | return self._set.__class__() # empty set 84 | 85 | startitem = self._set._item_by_index(start) 86 | stop -= start 87 | 88 | return self._set._from_items( 89 | islice(startitem._iter_to(None), 0, stop, step)) 90 | else: 91 | item = self._set._item_by_index(key) 92 | return item.key 93 | 94 | def __delitem__(self, key): 95 | if isinstance(key, slice): 96 | set_len = len(self._set) 97 | start, stop, step = key.indices(set_len) 98 | 99 | if step != 1: 100 | raise ValueError("Step is not suported for item deletion") 101 | if stop <= start: 102 | return # nothing to delete 103 | 104 | next, update = self._set._item_and_pointers_by_index(start) 105 | stop -= start 106 | for i in range(stop): 107 | item = next 108 | next = item[0].forward 109 | del self._set._mapping[item.key] 110 | self._set._delete_node(item, update) 111 | else: 112 | item, update = self._set._item_and_pointers_by_index(key) 113 | self._set._delete_node(item, update) 114 | 115 | 116 | class ScoreView: 117 | __slots__ = ('_set',) 118 | 119 | def __init__(self, set): 120 | self._set = set 121 | 122 | def __getitem__(self, key): 123 | if isinstance(key, slice): 124 | if key.step != None: 125 | raise ValueError("Step must be None") 126 | if(key.start is not None and key.stop is not None and 127 | key.start >= key.stop) or not len(self._set): 128 | return self._set.__class__() 129 | 130 | startitem = self._set._header[0].forward 131 | if key.start is not None: 132 | if startitem.score < key.start: 133 | startitem = self._set._item_by_score_left_incl(key.start) 134 | if startitem is None: 135 | # means key.start is greater than max score in set 136 | return self._set.__class__() 137 | 138 | if key.stop is None: 139 | return self._set._from_items(startitem._iter_to(None)) 140 | 141 | result = self._set.__class__() 142 | for item in startitem._iter_to(None): 143 | if item.score >= key.stop: 144 | break 145 | result[item.key] = item.score 146 | return result 147 | else: 148 | raise NotImplementedError('Only slicing by score supported') 149 | 150 | def __delitem__(self, key): 151 | if isinstance(key, slice): 152 | if key.step != None: 153 | raise ValueError("Step must be None") 154 | if(key.start is not None and key.stop is not None and 155 | key.start >= key.stop) or not len(self._set): 156 | return # nothing to delete 157 | 158 | if key.start is None: 159 | key.start = self._set._header[0].forward.score 160 | next, update = self._set._item_and_pointers_by_score_left_incl( 161 | key.start) 162 | if next is None: 163 | # means key.start is greater than max score in set 164 | return # nothing to delete 165 | 166 | if key.stop: 167 | while next and next.score <= key.stop: 168 | item = next 169 | next = item[0].forward 170 | del self._set._mapping[item.key] 171 | self._set._delete_node(item, update) 172 | else: 173 | while next: 174 | item = next 175 | next = item[0].forward 176 | del self._set._mapping[item.key] 177 | self._set._delete_node(item, update) 178 | else: 179 | raise NotImplementedError('Only slicing by score supported') 180 | 181 | 182 | def _random_level(): 183 | """Returns a random level for the new skiplist node 184 | 185 | The return value of this function is between 1 and ZSKIPLIST_MAXLEVEL 186 | (both inclusive), with a powerlaw-alike distribution where higher 187 | levels are less likely to be returned. 188 | """ 189 | level = 1 190 | while random.random() < 0.25: 191 | level += 1 192 | return level 193 | 194 | 195 | class SortedSet(MutableMapping): 196 | __slots__ = ('_level', '_mapping', '_header', '_tail', 197 | 'by_score', 'by_index') 198 | 199 | def __init__(self, source=None): 200 | self._level = 1 201 | self._mapping = {} 202 | self._header = Item(empty, empty) 203 | self._tail = None 204 | self.by_index = RankView(self) 205 | self.by_score = ScoreView(self) 206 | if source is not None: 207 | self.update(source) 208 | 209 | @classmethod 210 | def _from_items(cls, items): 211 | """Generates a new set from iterator over Item objects 212 | 213 | Usually used to make a new set from ``item._iter_to(other_item)`` 214 | """ 215 | self = cls() 216 | for item in items: 217 | self[item.key] = item.score 218 | return self 219 | 220 | def __iter__(self): 221 | start = self._header[0].forward # header is always empty 222 | if not start: 223 | return iter(()) 224 | return (item.key for item in start._iter_to(None)) 225 | 226 | def __reversed__(self): 227 | start = self._tail 228 | if not start: 229 | return iter(()) 230 | return (item.key for item in start._iter_backwards_to(None)) 231 | 232 | def __len__(self): 233 | return len(self._mapping) 234 | 235 | 236 | def __setitem__(self, key, score): 237 | if key in self: 238 | del self[key] 239 | # TODO probably optimize changing a value 240 | item = Item(key, score) 241 | self._mapping[key] = item 242 | 243 | rank = [None] * self._level 244 | update = [None] * self._level 245 | x = self._header 246 | for i in range(self._level-1, -1, -1): 247 | # store rank that is crossed to reach the insert position 248 | rank[i] = 0 if i == self._level-1 else rank[i+1] 249 | while x[i].forward and x[i].forward < item: 250 | rank[i] += x[i].span 251 | x = x[i].forward 252 | update[i] = x 253 | 254 | level = _random_level() 255 | if level > self._level: 256 | for i in range(self._level, level): 257 | assert len(rank) == i 258 | rank.append(0) 259 | assert len(update) == i 260 | update.append(self._header) 261 | update[i][i].span = len(self) 262 | self._level = level 263 | 264 | x = item 265 | for i in range(level): 266 | x[i].forward = update[i][i].forward 267 | update[i][i].forward = x 268 | 269 | # update span covered by update[i] as x is inserted here 270 | x[i].span = update[i][i].span - (rank[0] - rank[i]) 271 | update[i][i].span = (rank[0] - rank[i]) + 1 272 | 273 | # increment span for untouched levels 274 | for i in range(level, self._level): 275 | update[i][i].span += 1 276 | 277 | x.backward = None if update[0] == self._header else update[0] 278 | if x[0].forward: 279 | x[0].forward.backward = x 280 | else: 281 | self._tail = x 282 | 283 | def __getitem__(self, key): 284 | return self._mapping[key].score 285 | 286 | def __delitem__(self, key): 287 | item = self._mapping.pop(key) 288 | update = [None] * self._level 289 | 290 | x = self._header 291 | for i in range(self._level-1, -1, -1): 292 | while x[i].forward and x[i].forward < item: 293 | x = x[i].forward 294 | update[i] = x 295 | 296 | assert item == x[0].forward 297 | self._delete_node(item, update) 298 | 299 | def _delete_node(self, x, update): 300 | for i in range(self._level): 301 | if update[i][i].forward == x: 302 | update[i][i].span += x[i].span - 1 303 | update[i][i].forward = x[i].forward 304 | else: 305 | update[i][i].span -= 1 306 | assert update[i][i].span > 0 or not update[i][i].forward, update 307 | if x[0].forward: 308 | x[0].forward.backward = x.backward 309 | else: 310 | self._tail = x.backward 311 | while self._level > 1 and not self._header[self._level-1].forward: 312 | self._level -= 1 313 | 314 | def index(self, key): 315 | x = self._header 316 | rank = -1 # first key is always a header (Empty key) 317 | item = self._mapping[key] 318 | for i in range(self._level-1, -1, -1): 319 | while x[i].forward and x[i].forward <= item: 320 | rank += x[i].span 321 | x = x[i].forward 322 | if x.key == key: 323 | return rank 324 | raise KeyError(key, score) 325 | 326 | def _item_by_index(self, rank): 327 | if rank < 0: 328 | raise IndexError(rank) 329 | x = self._header 330 | traversed = -1 # first key is always a header (Empty key) 331 | for i in range(self._level-1, -1, -1): 332 | while x[i].forward and x[i].span + traversed <= rank: 333 | traversed += x[i].span 334 | x = x[i].forward 335 | if traversed == rank: 336 | return x 337 | raise IndexError(rank) 338 | 339 | def _item_and_pointers_by_index(self, rank): 340 | if rank < 0: 341 | raise IndexError(rank) 342 | x = self._header 343 | update = [None] * self._level 344 | traversed = -1 # first key is always a header (Empty key) 345 | for i in range(self._level-1, -1, -1): 346 | while x[i].forward and x[i].span + traversed < rank: 347 | traversed += x[i].span 348 | x = x[i].forward 349 | update[i] = x 350 | x = x[0].forward 351 | return x, update 352 | 353 | def _item_by_score_left_incl(self, score): 354 | """Returns left most item scored up to `score` inclusive""" 355 | # "left" name is analogy to bisect_left 356 | x = self._header 357 | for i in range(self._level-1, -1, -1): 358 | while x[i].forward and x[i].forward.score < score: 359 | x = x[i].forward 360 | x = x[0].forward 361 | return x 362 | 363 | def _item_and_pointers_by_score_left_incl(self, score): 364 | """Returns left most item scored up to `score` inclusive 365 | 366 | This one returns also ``update`` array to assist in item deletion 367 | """ 368 | x = self._header 369 | update = [None] * self._level 370 | for i in range(self._level-1, -1, -1): 371 | while x[i].forward and x[i].forward.score < score: 372 | x = x[i].forward 373 | update[i] = x 374 | x = x[0].forward 375 | return x, update 376 | 377 | def __repr__(self): 378 | return ''.format(reprlib.Repr().repr_dict(self, 1)) 379 | 380 | 381 | -------------------------------------------------------------------------------- /test_sortedsets.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import random 3 | import copy 4 | import fractions 5 | from operator import itemgetter 6 | from itertools import combinations, product 7 | from unittest.mock import patch 8 | 9 | from sortedsets import SortedSet 10 | 11 | 12 | class TestSortedSets(unittest.TestCase): 13 | 14 | def test_simple(self): 15 | ss = SortedSet() 16 | ss['one'] = 1 17 | ss['two'] = 2 18 | self.assertEqual(ss['one'], 1) 19 | self.assertEqual(ss['two'], 2) 20 | 21 | def test_index(self): 22 | ss = SortedSet({ 23 | 'one': 1, 24 | 'two': 2, 25 | }) 26 | self.assertEqual(ss.index('two'), 1) 27 | self.assertEqual(ss.index('one'), 0) 28 | 29 | def test_keys(self): 30 | ss = SortedSet() 31 | ss['one'] = 1 32 | ss['two'] = 2 33 | self.assertEqual(list(ss), [ 34 | 'one', 35 | 'two', 36 | ]) 37 | self.assertEqual(list(ss.keys()), [ 38 | 'one', 39 | 'two', 40 | ]) 41 | 42 | def test_items(self): 43 | ss = SortedSet() 44 | ss['one'] = 1 45 | ss['two'] = 2 46 | self.assertEqual(list(ss.items()), [ 47 | ('one', 1), 48 | ('two', 2), 49 | ]) 50 | 51 | def test_values(self): 52 | ss = SortedSet() 53 | ss['one'] = 1 54 | ss['two'] = 2 55 | self.assertEqual(list(ss.values()), [1, 2]) 56 | 57 | def test_negative(self): 58 | ss = SortedSet() 59 | ss['one'] = 1 60 | ss['two'] = -2 61 | self.assertEqual(ss['one'], 1) 62 | self.assertEqual(ss['two'], -2) 63 | self.assertEqual(list(ss), ['two', 'one']) 64 | self.assertEqual(ss.index('two'), 0) 65 | self.assertEqual(ss.index('one'), 1) 66 | 67 | def test_floats(self): 68 | ss = SortedSet() 69 | # we use values that have exact representation as floating point number 70 | ss['one'] = 1.25 71 | ss['two'] = 1.5 72 | ss['three'] = -3.0 73 | self.assertEqual(ss['one'], 1.25) 74 | self.assertEqual(ss['two'], 1.5) 75 | self.assertEqual(ss['three'], -3.0) 76 | self.assertEqual(list(ss), ['three', 'one', 'two']) 77 | 78 | def test_fractions(self): 79 | ss = SortedSet() 80 | # we use values that have exact representation as floating point number 81 | ss['one'] = fractions.Fraction(1, 2) 82 | ss['two'] = fractions.Fraction(2, 3) 83 | ss['three'] = fractions.Fraction(-3, 2) 84 | self.assertEqual(ss['one'], fractions.Fraction(1, 2)) 85 | self.assertEqual(ss['two'], fractions.Fraction(2, 3)) 86 | self.assertEqual(ss['three'], fractions.Fraction(-3, 2)) 87 | self.assertEqual(list(ss), ['three', 'one', 'two']) 88 | 89 | def test_delete(self): 90 | ss = SortedSet() 91 | ss['one'] = 1 92 | ss['two'] = 2 93 | ss['three'] = 3 94 | del ss['two'] 95 | self.assertEqual(ss['one'], 1) 96 | self.assertEqual(ss['three'], 3) 97 | self.assertEqual(list(ss), ['one', 'three']) 98 | 99 | def test_key_by_index(self): 100 | ss = SortedSet({ 101 | 'one': 1, 102 | 'two': 2, 103 | }) 104 | self.assertEqual(ss.by_index[1], 'two') 105 | self.assertEqual(ss.by_index[0], 'one') 106 | 107 | def test_repr(self): 108 | ss = SortedSet({ 109 | 'one': 1, 110 | 'two': 2, 111 | }) 112 | self.assertEqual(repr(ss), "") 113 | 114 | def test_slice_by_index(self): 115 | data = { 116 | 'one': 1, 117 | 'two': 2, 118 | } 119 | ss = SortedSet(data) 120 | self.assertEqual(ss.by_index[:], SortedSet(data)) 121 | self.assertEqual(ss.by_index[1:], SortedSet({'two': 2})) 122 | self.assertEqual(ss.by_index[2:], SortedSet()) 123 | self.assertEqual(ss.by_index[3:], SortedSet()) 124 | self.assertEqual(ss.by_index[:-1], SortedSet({'one': 1})) 125 | self.assertEqual(ss.by_index[:-2], SortedSet()) 126 | self.assertEqual(ss.by_index[:-3], SortedSet()) 127 | self.assertEqual(ss.by_index[:-4], SortedSet()) 128 | with self.assertRaises(ValueError): 129 | ss.by_index[::0] 130 | with self.assertRaises(ValueError): 131 | ss.by_index[::-1] 132 | with self.assertRaises(ValueError): 133 | ss.by_index[::-2] 134 | ss = SortedSet(data) 135 | del ss.by_index[:] 136 | self.assertEqual(ss, SortedSet()) 137 | ss = SortedSet(data) 138 | del ss.by_index[1:] 139 | self.assertEqual(ss, SortedSet({'one': 1})) 140 | ss = SortedSet(data) 141 | del ss.by_index[:-1] 142 | self.assertEqual(ss, SortedSet({'two': 2})) 143 | 144 | def test_slice_by_score(self): 145 | data = { 146 | 'one': 1, 147 | 'two1': 2, 148 | 'two2': 2, 149 | 'two3': 2, 150 | 'three': 3, 151 | } 152 | ss = SortedSet(data) 153 | self.assertEqual(ss.by_score[:], SortedSet(data)) 154 | self.assertEqual(ss.by_score[1:], SortedSet(data)) 155 | self.assertEqual(ss.by_score[1:4], SortedSet(data)) 156 | self.assertEqual(ss.by_score[:4], SortedSet(data)) 157 | tmp = data.copy() 158 | del tmp['one'] 159 | self.assertEqual(ss.by_score[1.5:], SortedSet(tmp)) 160 | self.assertEqual(ss.by_score[2:], SortedSet(tmp)) 161 | self.assertEqual(ss.by_score[2:4], SortedSet(tmp)) 162 | self.assertEqual(ss.by_score[1.5:4], SortedSet(tmp)) 163 | self.assertEqual(ss.by_score[1.5:3.1], SortedSet(tmp)) 164 | tmp = data.copy() 165 | del tmp['three'] 166 | self.assertEqual(ss.by_score[0.5:3], SortedSet(tmp)) 167 | self.assertEqual(ss.by_score[0.5:2.5], SortedSet(tmp)) 168 | self.assertEqual(ss.by_score[:2.5], SortedSet(tmp)) 169 | self.assertEqual(ss.by_score[:3], SortedSet(tmp)) 170 | self.assertEqual(ss.by_score[:2.01], SortedSet(tmp)) 171 | del tmp['one'] 172 | self.assertEqual(ss.by_score[1.01:2.01], SortedSet(tmp)) 173 | self.assertEqual(ss.by_score[1.99:2.01], SortedSet(tmp)) 174 | self.assertEqual(ss.by_score[2:2.01], SortedSet(tmp)) 175 | self.assertEqual(ss.by_score[2:2.99], SortedSet(tmp)) 176 | self.assertEqual(ss.by_score[2.01:2.99], SortedSet()) 177 | self.assertEqual(ss.by_score[:2], SortedSet({'one': 1})) 178 | self.assertEqual(ss.by_score[:1.01], SortedSet({'one': 1})) 179 | self.assertEqual(ss.by_score[:1.99], SortedSet({'one': 1})) 180 | self.assertEqual(ss.by_score[3:], SortedSet({'three': 3})) 181 | self.assertEqual(ss.by_score[2.01:], SortedSet({'three': 3})) 182 | self.assertEqual(ss.by_score[2.99:], SortedSet({'three': 3})) 183 | self.assertEqual(ss.by_score[3.01:], SortedSet()) 184 | 185 | self.assertEqual(ss.by_score[ 186 | fractions.Fraction(4/3):fractions.Fraction(7/3) 187 | ], SortedSet({'two1': 2, 'two2': 2, 'two3': 2})) 188 | 189 | def test_delete_all_cases(self): 190 | for levels in product(range(1, 4), range(1, 4), range(1, 4)): 191 | # delete middle 192 | ss = SortedSet() 193 | with patch('sortedsets._random_level', side_effect=levels) as p: 194 | ss['one'] = 1 195 | ss['two'] = 2 196 | ss['three'] = 3 197 | del ss['two'] 198 | self.assertEqual(ss['one'], 1) 199 | self.assertEqual(ss['three'], 3) 200 | with self.assertRaises(KeyError): 201 | ss['two'] 202 | self.assertEqual(list(ss), ['one', 'three']) 203 | ss.clear() 204 | self.assertEqual(list(ss), []) 205 | # delete first 206 | ss = SortedSet() 207 | with patch('sortedsets._random_level', side_effect=levels) as p: 208 | ss['one'] = 1 209 | ss['two'] = 2 210 | ss['three'] = 3 211 | del ss['one'] 212 | self.assertEqual(ss['two'], 2) 213 | self.assertEqual(ss['three'], 3) 214 | with self.assertRaises(KeyError): 215 | ss['one'] 216 | self.assertEqual(list(ss), ['two', 'three']) 217 | ss.clear() 218 | self.assertEqual(list(ss), []) 219 | # delete last 220 | ss = SortedSet() 221 | with patch('sortedsets._random_level', side_effect=levels) as p: 222 | ss['one'] = 1 223 | ss['two'] = 2 224 | ss['three'] = 3 225 | del ss['three'] 226 | self.assertEqual(ss['one'], 1) 227 | self.assertEqual(ss['two'], 2) 228 | with self.assertRaises(KeyError): 229 | ss['three'] 230 | self.assertEqual(list(ss), ['one', 'two']) 231 | ss.clear() 232 | self.assertEqual(list(ss), []) 233 | 234 | 235 | class TestFuzzy(unittest.TestCase): 236 | 237 | def test_insert_integers(self): 238 | items = [ # fifty random values 239 | ('Xe2W0QxllGdCW251l7U9Dg', 150), 240 | ('3HT/SVSdCwM+4ZjtSqHCew', 476), 241 | ('Q2BKuEOFwIojkPsnjmPNFg', 2390), 242 | ('Qjq1fGnHVc5nXvWnbEyXjQ', 3773), 243 | ('wamjIdfm+ajk81fR7gKcAA', 4729), 244 | ('uDUHFv5CtiY/Gm5LOCfGUg', 6143), 245 | ('78l934GXETN68sql2vjP5w', 6487), 246 | ('1LAwgvIO0tEikYySkaSaXw', 7449), 247 | ('Y4mDqz7LfIV4L8h2aUwAkA', 10234), 248 | ('AidPceym19y/lmIdi6JxQQ', 10779), 249 | ('hHqwNSMusq7O895yFkr+rQ', 10789), 250 | ('dg16QiUDC2rgE39FWTSOxg', 11873), 251 | ('sgAdgtQ5wRFGSOZ3xZYHyA', 12273), 252 | ('OACKY0A1ftBbyLvTzyf8lQ', 12776), 253 | ('f+dLA1jK8EFEAHxm1FKUkA', 13079), 254 | ('1uDN4mSmsEQF/o6VNiBl3g', 13147), 255 | ('nNwOvGfk9AH2tIzK8uNdzQ', 16636), 256 | ('tMUZ6A1e/1SKd3ko0FhhBQ', 17933), 257 | ('M77ZQiFlYeU4ySUtVa6XYg', 18570), 258 | ('fY7RKQu8toBxoug4CMmznA', 18711), 259 | ('UaZorA+/GnCL4nmgLs3c/g', 19968), 260 | ('VbXaOsRHqH2CAoNJiYsrqg', 20064), 261 | ('dAr84/axpItIAjjNcVPzHA', 20250), 262 | ('HjzS0QlpofFhDO2iU4mXAw', 20756), 263 | ('ipksmQaeYErYwjZ6ia46TA', 21084), 264 | ('XemDsutAYPpCZ6WY4M//ig', 21609), 265 | ('6U6fbOs8jYVfqWeArQ5HHQ', 22410), 266 | ('QFblGefWYZqFbuSK0SDPRQ', 23267), 267 | ('J13bR75czCiLfeBcIy4XTg', 26556), 268 | ('e6XlDT9h6GVPdfvBOrXW5Q', 26608), 269 | ('/eLYo+GKgAt7I2PrOrFTzQ', 28597), 270 | ('48W/xF+VIQZoyKlhktifMw', 31364), 271 | ('NTUtbi4YOHiNIV6SVrpwyg', 33709), 272 | ('364+KUYYuwlmezv1EvaE0g', 33945), 273 | ('YaD6Ktqw1iIWcFGgMEvMxg', 38248), 274 | ('cJSZfsidFuaMK9jY15g44A', 38668), 275 | ('UeP/HvscsnQXUK37Dyo8/w', 40861), 276 | ('xon2bN9ZToI4LpN4o7M2nQ', 41836), 277 | ('MQKXJCNNtWsRqoGbSaDorw', 47171), 278 | ('LCcqUwfmOFq+VXI2kGBQow', 49311), 279 | ('gMXF4DMHCWBjbgucOqWKQg', 50725), 280 | ('JKHDvGMcLQrR4G3zC2g9ug', 50875), 281 | ('Mp1feZZmnmMPJk8bGv0NaA', 51017), 282 | ('rhZyspOoakQBO9Ses3jl+A', 53781), 283 | ('JB9bMHKHoT+hMVjuBrbqlg', 56409), 284 | ('/DsgGH+7F6Fh2/81SzyXYA', 56512), 285 | ('InjjAuUMGHYUIRdRnkUw2w', 56903), 286 | ('otVFi6DLAO+v7XUAcmKttA', 57114), 287 | ('mVTvHObgjfzvZLOzl/xo2Q', 58550), 288 | ('uU1yLoXCgPtifROhCST0sA', 60267)] 289 | keys = list(map(itemgetter(0), items)) 290 | values = list(map(itemgetter(1), items)) 291 | 292 | # simple checks 293 | set_n = SortedSet(items) 294 | set_r = SortedSet(reversed(items)) 295 | for cur in (set_n, set_r): 296 | self.assertEqual(list(cur), keys) 297 | self.assertEqual(list(cur.keys()), keys) 298 | self.assertEqual(list(cur.values()), values) 299 | self.assertEqual(list(cur.items()), items) 300 | for idx in range(len(items)): 301 | key, score = items[idx] 302 | self.assertEqual(set_n.index(key), idx) 303 | self.assertEqual(set_r.index(key), idx) 304 | self.assertEqual(set_n.by_index[idx], key) 305 | self.assertEqual(set_r.by_index[idx], key) 306 | 307 | # slicing by index should be same as slicing list 308 | ends = (None, 0, 2, 5, 10, 49, 50, 51, -1, -2, -5, -10, -49, 50, -51) 309 | steps = (1, 2, 3) 310 | for start, stop, step in product(ends, ends, steps): 311 | self.assertEqual(list(set_n.by_index[start:stop:step].items()), 312 | items[start:stop:step]) 313 | self.assertEqual(list(set_r.by_index[start:stop:step].items()), 314 | items[start:stop:step]) 315 | # let's try to delete and reinsert slice 316 | for start, stop in product(ends, ends): 317 | slc = set_r.by_index[start:stop] 318 | slclist = items[start:stop] 319 | self.assertEqual(list(slc.items()), slclist) 320 | # delete a slice from set and list, then compare 321 | del set_r.by_index[start:stop] 322 | tmp = items[:] 323 | del tmp[start:stop] 324 | self.assertEqual(list(set_r.items()), tmp) 325 | # let's recreate set and check 326 | set_r.update(slc) 327 | self.assertEqual(list(set_r.items()), items) 328 | 329 | # try score slice on reverted set 330 | if slclist: 331 | # can do this only if set is not empty 332 | scorestart = slclist[0][1] 333 | scorestop = slclist[-1][1] + 1 334 | slc2 = set_r.by_score[scorestart:scorestop] 335 | self.assertEqual(list(slc2.items()), slclist) 336 | del set_r.by_score[scorestart:scorestop] 337 | self.assertEqual(list(set_r.items()), tmp) 338 | # and revert this change back again 339 | set_r.update(slc2) 340 | self.assertEqual(list(set_r.items()), items) 341 | 342 | 343 | # Lets test 7 random insertion orders 344 | for i in (10, 20, 30, 40, 50, 60, 70): 345 | rnd = random.Random(i) 346 | to_insert = copy.copy(items) 347 | rnd.shuffle(to_insert) 348 | 349 | # Let's check several sets created with different methods 350 | set1 = SortedSet() 351 | for k, v in to_insert: 352 | set1[k] = v 353 | set2 = SortedSet(to_insert) 354 | set3 = SortedSet(set1) 355 | set4 = SortedSet(set2) 356 | 357 | # Check all of them 358 | cursets = (set1, set2, set3, set4) 359 | for cur in cursets: 360 | self.assertEqual(list(cur), keys) 361 | self.assertEqual(list(cur.keys()), keys) 362 | self.assertEqual(list(cur.values()), values) 363 | self.assertEqual(list(cur.items()), items) 364 | 365 | # Check equality of all combinations 366 | all_sets = (set_n, set_r, set1, set2, set3, set4) 367 | for s1, s2 in combinations(all_sets, 2): 368 | self.assertEqual(s1, s2) 369 | 370 | # Let's pick up items to delete 371 | left = copy.copy(items) 372 | to_delete = [] 373 | for i in range(rnd.randrange(10, 30)): 374 | idx = random.randrange(len(left)) 375 | to_delete.append(left[idx]) 376 | del left[idx] 377 | 378 | # Let's test deletion 379 | for cur in cursets: 380 | rnd.shuffle(to_delete) 381 | for key, value in to_delete: 382 | del cur[key] 383 | self.assertEqual(list(cur.items()), left) 384 | self.assertNotEqual(cur, set_n) 385 | self.assertNotEqual(cur, set_r) 386 | 387 | # Let's reinsert keys in random order, and check if it's still ok 388 | for cur in cursets: 389 | rnd.shuffle(to_delete) 390 | for key, value in to_delete: 391 | cur[key] = value 392 | self.assertEqual(cur, set_n) 393 | self.assertEqual(cur, set_r) 394 | self.assertEqual(list(cur.items()), items) 395 | for idx in range(len(items)): 396 | key, score = items[idx] 397 | self.assertEqual(cur.index(key), idx) 398 | self.assertEqual(cur.by_index[idx], key) 399 | 400 | 401 | if __name__ == '__main__': 402 | unittest.main() 403 | --------------------------------------------------------------------------------