import unittest
import doctest
try:
    import rbtree
except ImportError:
    # we are running it from the build directory
    import testpath
    testpath.modTestPath()
    import rbtree

def until(rbt, from_iter, stopkey):
    target = iter(rbt)
    target.goto(stopkey)
    yield from_iter.key
    # We don't worry about lte/gte/direction because we are testing the
    # positioned iter
    while from_iter != target:
        yield from_iter.next()

def at(rbt, key):
    i = iter(rbt)
    i.goto(key)
    return i


class Test(unittest.TestCase):
    def supply_words(self, ds):
        for word in open('/usr/share/dict/words', 'r'):
            word = word.strip()
            ds[word] = True

    def sample_data(self):
        r = rbtree.RBTree()
        self.supply_words(r)
        return r

    def test_create(self):
        r = rbtree.RBTree()
        assert len(r) == 0

    def test_create_mapping(self):
        r = rbtree.RBTree({'a' : 'b'})
        assert len(r) == 1
        assert r['a'] == 'b'

    def test_dict(self):
        r = rbtree.RBTree({'a' : 'b'})
        d = dict(r)
        assert d['a'] == 'b'

    def test_create_sequence(self):
        r = rbtree.RBTree((('a', 'b'), ('c', 'd')))
        assert len(r) == 2
        assert r['a'] == 'b'
        assert r['c'] == 'd'

    def test_missing(self):
        r = rbtree.RBTree({'a' : 'b'})
        assert r['a']
        try:
            print r['missing']
        except KeyError:
            pass
        else:
            raise KeyError("Able to find missing key")

    def test_iteration(self):
        r = self.sample_data()
        i1 = iter(r)
        i2 = iter(r)
        for word in i1:
            assert i2.next() == word


    def test_axes(self):
        r = rbtree.RBTree(dict(zip(range(10), range(1, 11))))
        assert r.keys() == range(10)
        assert r.values() == range(1, 11)
        assert r.items() == zip(range(10), range(1, 11))

    def test_slicing(self):
        r = self.sample_data()
        # This should return all matches between b and c in the wordlist
        slice = r['b':'c']
        assert len(slice)
        assert slice.keys()
        del slice
        # Empty slices return empty trees
        r = rbtree.RBTree()
        assert len(r[0:100]) == 0

    def test_slicing_stepping(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        assert r[0:8:2].keys() == [0,2,4,6]
        assert r[1:8:3].keys() == [1,4,7]
        # neg slice stride makes no sense with cmp ordered keys
        assert r[0:8:-2].keys() == [0,2,4,6]

    def test_slice_partial(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        assert r[8:]


    def test_iter_refcount(self):
        r = rbtree.RBTree(dict(zip(range(10), range(1, 11))))
        i = iter(r)
        del r
        assert list(i) == range(10)

    def test_pop(self):
        r = rbtree.RBTree(dict(zip(range(10), range(1, 11))))
        assert len(r) == 10
        # Pops the values, not the keys
        assert r.pop() == 1
        assert r.pop() == 2
        assert len(r) == 8


    def test_random_insert(self):
        import random
        r = rbtree.RBTree()
        nums = range(1000)
        random.shuffle(nums)
        for i in nums:
            r[i] = True

        assert r.keys() == range(1000)
        assert r[800:810].keys() == range(800, 810)

    def test_contains(self):
        r = rbtree.RBTree(dict(zip(range(10), range(1, 11))))
        assert 5 in r
        assert r.has_key(5)
        assert 42 not in r

    def test_copy(self):
        r = rbtree.RBTree(dict(zip(range(10), range(1, 11))))
        c = r.copy()
        assert 5 in r
        assert 5 in c
        del r[5]
        assert 5 not in r
        assert 5 in c

    def test_clear(self):
        r = rbtree.RBTree(dict(zip(range(10), range(1, 11))))
        r.clear()
        assert len(r) == 0
        assert 5 not in r

    def test_setdefault(self):
        r = rbtree.RBTree()
        d = r.setdefault('f', 'foo')
        assert r['f'] == 'foo'

    def test_cmp(self):
        reverse = lambda x, y: cmp(x,y)*-1
        r = rbtree.RBTree(dict(zip(range(10), range(10))),
                          cmp=reverse)
        assert r.keys() == [9,8,7,6,5,4,3,2,1,0]

        badcmp = lambda x,y: 'z'
        try:
            r = rbtree.RBTree(dict(zip(range(10), range(10))),
                              cmp=badcmp)
        except TypeError: pass
        else: raise TypeError("Allowed Cmp that returns string?")

    def test_cmp_slice(self):
        # make sure that the slice keeps a reference to the compare
        # method
        # this is with the intention that multikey rbtrees do the
        # right thing
        def multi(a, b):
            if a == b: return 1
            return a - b
        r = rbtree.RBTree(dict(zip(range(10), range(10))), cmp=multi)
        r.update(dict(zip(range(10), range(10))))
        # there are now two of each set
        assert len(r) == 20
        s = r[:3]
        assert len(s) == 8
        # if the len was 4 then the default key compare would have
        # collapsed the rbtree during construction

    def test_minmax(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        assert r.min() == 0
        assert r.max() == 9


    def test_offset(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        # get keys by offset
        assert r.byOffset(2) == 2
        assert r.byOffset(1) == 1
        assert r.byOffset(0) == 0
        assert r.byOffset(-1) == 9
        assert r.byOffset(-2) == 8

    def test_stopiteration(self):
        r = rbtree.RBTree({'a': 'b', 'c': 'd'})
        i = iter(r)
        # iterate the whole set
        list(i)
        # now each call should raise StopIteration to complie with the
        # protocol
        try:
            i.next()
        except StopIteration:
            pass
        else:
            raise NotImplementedError()

        try:
            i.next()
        except StopIteration:
            pass
        else:
            raise NotImplementedError()

    def test_iterdirection(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        i = iter(r)
        assert i.next() == 0
        assert i.next() == 1
        assert i.next() == 2
        i.direction = -1
        assert i.next() == 1
        assert i.next() == 0

        i = iter(r)
        i.direction = -1
        assert i.next() == 9

    def test_itergoto(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        i = iter(r)
        i.goto(5)
        assert i.next() == 6
        assert i.next() == 7

        # Just to key 5 and walk backwards getting the items tuple
        i2 = r.iteritems()
        i2.goto(5)
        i2.direction = -1
        assert i2.next() == (4,4)

        # Now goto a missing key
        try:
            i2.goto('missing')
        except KeyError:
            pass
        else:
            raise KeyError("Found Missing key in iter")

    def test_iternear(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        i = iter(r)
        i.goto(4.5, near=True)
        assert i.key == 5

        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        i = iter(r)
        i.goto(4.5, near=True)
        assert i.key == 5
        z = iter(r)
        while True:
            z.next()
            if z == i: break
        assert z.key == 5

    def test_itercmp(self):
        r = rbtree.RBTree(dict(zip(range(10), range(10))))
        x = iter(r).goto(7)
        assert list(until(r, at(r, 5), 7)) == [5,6,7]

    def test_replace(self):
        r = rbtree.RBTree()
        r[1] = 2
        assert r[1] == 2
        r[1] = 3
        assert r[1] == 3
        assert len(r) == 1

    def test_multikey_compare(self):
        # use a cmp function that allows for key matches
        def multicmp(x, y):
            x = cmp(x, y)
            if x == 0: return 1
            return x
        r = rbtree.RBTree(cmp=multicmp)
        r[1] = 2
        r[1] = 3
        assert len(r) == 2
        assert 2 in r.values()
        assert 3 in r.values()

        i = iter(r)
        i.next()
        i.delete()
        assert len(r) == 1

    def test_iterprev(self):
        r = rbtree.RBTree([(chr(i), True) for i in range(ord('a'),
                                                         ord('z')+1)])
        i = r.iternodes()
        i.goto('c')
        assert i.key == 'c'
        i.prev()
        assert i.key == 'b'
        i.next()
        assert i.key == 'c'

        # Now we flip the base direction
        i.direction = -1
        i.next()
        assert i.key == 'b'
        i.prev()
        assert i.key == 'c'

    def test_pickle(self):
        import pickle
        r = rbtree.RBTree([(chr(i), True) for i in range(ord('a'),
                                                         ord('z')+1)])
        s = pickle.dumps(r, protocol=-1)
        t = pickle.loads(s)
        assert t.keys() == [chr(i) for i in range(ord('a'), ord('z')+1)]

        # and with a custom cmp function, this has to be module global
        # to work because of pickle
        global ncmp
        def ncmp(a, b):
            return cmp(a.lower(), b.lower())
        r = rbtree.RBTree({'a' : 1, 'A' : 2}, ncmp)
        s = pickle.dumps(r, protocol=-1)
        t = pickle.loads(s)
        assert t == r


    def test_slice_endcase(self):
        r = rbtree.RBTree((('x', 1), ('y', 1), ('z', 1)))
        assert r['y':].keys() == ['y', 'z']

    def test_delslice(self):
        r = rbtree.RBTree([(chr(i), True) for i in range(ord('a'),
                                                         ord('z')+1)])
        del r ['a':'x']
        assert r.keys() == ['x', 'y', 'z']
        del r['y':]
        assert r.keys() == ['x']

    def test_iteroffset(self):
        r = rbtree.RBTree([(chr(i), True) for i in range(ord('a'),
                                                         ord('z')+1)])
        i = iter(r)
        assert i.byOffset(-1).key == 'z'
        assert i.byOffset(0).key  == 'a'
        assert i.byOffset(12).key == 'm'

if __name__ == "__main__":
    suite = unittest.TestSuite()
    suite.addTest(unittest.makeSuite(Test))
    suite.addTest(doctest.DocFileSuite('README.txt'))
    runner = unittest.TextTestRunner(verbosity=1)
    runner.run(suite)

