23 April 2007

N-based Python list

Solving http://www.answermysearches.com/python-subclassing-list-to-make-it-n-based/251/ (with some light testing)
class n_based_list(list):
    def __init__(self, startidx, items):
        self.start = startidx
        list.__init__(self, items)

    def idx(self, i):
        return i - self.start

    def insert(self, i, x):
        return list.insert(self, self.idx(i), x)

    def pop(self, i):
        return list.pop(self, self.idx(i))

    def __getslice__(self, i, j):
        return list.__getslice__(self, self.idx(i), j)

    def _handle_slice(self, s):
        return slice(self.idx(s.start or self.start),
                     s.stop,
                     s.step)

    def __getitem__(self, i):
        if isinstance(i, slice):
            return list.__getitem__(self, self._handle_slice(i))
        return list.__getitem__(self, self.idx(i))

    def __setitem__(self, i, x):
        if isinstance(i, slice):
            return list.__setitem__(self, self.handle_slice(i), x)
        return list.__setitem__(self, self.idx(i), x)

    def __delitem__(self, i):
        if isinstance(i, slice):
            return list.__delitem__(self, self._handle_slice(i))
        return list.__delitem__(self, self.idx(i))
    
    def enumerate(self):
        for i, elt in enumerate(self):
            yield i+self.start, elt
            
def test_n_based_list():
    l = n_based_list(1, ['a', 'b', 'c', 'd'])
    assert l[1] == 'a'
    assert l[0] == 'd' # yuck!
    assert l[2] == 'b'
    assert l[3] == 'c'
    assert l[4] == 'd'

    assert l.pop(1) == 'a'
    assert l == ['b', 'c', 'd']
    l.insert(1, 'a')
    assert l[1] == 'a'

    assert l[1:3] == ['a', 'b', 'c']
    assert l[1::2] == ['a', 'c']
    del l[::2]
    assert l == ['b', 'd']
    new_l = n_based_list(1, ['a', 'b', 'c', 'd'])
    assert list(new_l.enumerate()) == [(1, 'd'),
                                       (2, 'a'),
                                       (3, 'b'),
                                       (4, 'c')]