Skip to content
This repository
Browse code

Merge pull request #212 from lukecampbell/orderedset

Adds OrderedSet to pyon utils
  • Loading branch information...
commit e0e3ab514c4335d88cab7fd47d2ac65ce8892d87 2 parents 5841d02 + 2cac450
Dave Foster daf authored

Showing 2 changed files with 133 additions and 0 deletions. Show diff stats Hide diff stats

  1. +71 0 pyon/util/set.py
  2. +62 0 pyon/util/test/test_set.py
71 pyon/util/set.py
... ... @@ -0,0 +1,71 @@
  1 +#!/usr/bin/env python
  2 +
  3 +'''
  4 +@author (http://code.activestate.com/recipes/576694/)
  5 +@file pyon/util/set.py
  6 +@date Wed Jan 2 10:43:06 EST 2013
  7 +'''
  8 +## {{{ http://code.activestate.com/recipes/576694/ (r9)
  9 +import collections
  10 +
  11 +class OrderedSet(collections.MutableSet):
  12 +
  13 + def __init__(self, iterable=None):
  14 + self.end = end = []
  15 + end += [None, end, end] # sentinel node for doubly linked list
  16 + self.map = {} # key --> [key, prev, next]
  17 + if iterable is not None:
  18 + self |= iterable
  19 +
  20 + def __len__(self):
  21 + return len(self.map)
  22 +
  23 + def __contains__(self, key):
  24 + return key in self.map
  25 +
  26 + def add(self, key):
  27 + if key not in self.map:
  28 + end = self.end
  29 + curr = end[1]
  30 + curr[2] = end[1] = self.map[key] = [key, curr, end]
  31 +
  32 + def discard(self, key):
  33 + if key in self.map:
  34 + key, prev, next = self.map.pop(key)
  35 + prev[2] = next
  36 + next[1] = prev
  37 +
  38 + def __iter__(self):
  39 + end = self.end
  40 + curr = end[2]
  41 + while curr is not end:
  42 + yield curr[0]
  43 + curr = curr[2]
  44 +
  45 + def __reversed__(self):
  46 + end = self.end
  47 + curr = end[1]
  48 + while curr is not end:
  49 + yield curr[0]
  50 + curr = curr[1]
  51 +
  52 + def pop(self, last=True):
  53 + if not self:
  54 + raise KeyError('set is empty')
  55 + key = self.end[1][0] if last else self.end[2][0]
  56 + self.discard(key)
  57 + return key
  58 +
  59 + def __repr__(self): #pragma no cover
  60 + if not self:
  61 + return '%s()' % (self.__class__.__name__,)
  62 + return '%s(%r)' % (self.__class__.__name__, list(self))
  63 +
  64 + def __eq__(self, other):
  65 + if isinstance(other, OrderedSet):
  66 + return len(self) == len(other) and list(self) == list(other)
  67 + return set(self) == set(other)
  68 +
  69 +
  70 +## end of http://code.activestate.com/recipes/576694/ }}}
  71 +
62 pyon/util/test/test_set.py
... ... @@ -0,0 +1,62 @@
  1 +#!/usr/bin/env/python
  2 +
  3 +'''
  4 +@author Luke Campbell <LCampbell (at) ASAScience.com>
  5 +@file pyon/util/test/test_set.py
  6 +@date Wed Jan 2 10:46:48 EST 2013
  7 +'''
  8 +
  9 +from pyon.util.set import OrderedSet
  10 +from pyon.util.unit_test import PyonTestCase
  11 +from nose.plugins.attrib import attr
  12 +from copy import copy
  13 +
  14 +@attr('UNIT')
  15 +class OrderedSetTest(PyonTestCase):
  16 + def test_set_ops(self):
  17 + # Verify unique keys and order is maintained
  18 + os = OrderedSet([3,1,5,2,2,1,9,8,12,18])
  19 + compare_list = [3,1,5,2,9,8,12,18]
  20 + self.assertEquals(list(os), compare_list)
  21 +
  22 +
  23 + # Verify addition doesn't violate unique keys
  24 + os.add(2)
  25 + self.assertEquals(list(os), compare_list)
  26 +
  27 + # Verify a valid addition doesn't violate order
  28 + os.add(21)
  29 + self.assertEquals(list(os), compare_list + [21])
  30 +
  31 + # Pop
  32 + val = os.pop()
  33 + self.assertEquals(val, 21)
  34 + self.assertEquals(list(os), compare_list)
  35 +
  36 + # Verify list comprehension doesn't violate order
  37 + self.assertEquals([i for i in os], compare_list)
  38 +
  39 + # Lenght assertions
  40 + self.assertEquals(len(list(os)), len(os))
  41 +
  42 + # Reversed
  43 + cmpval = copy(compare_list)
  44 + cmpval.reverse()
  45 + self.assertEquals([i for i in reversed(os)], cmpval)
  46 +
  47 +
  48 + # Comparisons
  49 + self.assertTrue(os == compare_list)
  50 + os2 = OrderedSet(compare_list)
  51 + self.assertTrue(os == os2)
  52 +
  53 +
  54 + # Discarding
  55 + self.assertTrue(8 in os)
  56 + os.discard(8)
  57 + self.assertFalse(8 in os)
  58 +
  59 +
  60 + with self.assertRaises(KeyError):
  61 + os2 = OrderedSet([])
  62 + os2.pop()

0 comments on commit e0e3ab5

Please sign in to comment.
Something went wrong with that request. Please try again.