diff --git a/aaf/iterator.pyx b/aaf/iterator.pyx index 940fb3e..2f4973b 100644 --- a/aaf/iterator.pyx +++ b/aaf/iterator.pyx @@ -11,12 +11,32 @@ cdef class BaseIterator(object): def __cinit__(self): self._clone_iter = None - def __getitem__(self, int index): + def __getitem__(self, index): + + if isinstance(index, slice): + return self._getslice(index) + + index = int(index) + + if index < 0: + index = len(self) + index + + if index < 0: + raise IndexError("index out of range") + for i, item in enumerate(self): if i == index: return item raise IndexError("index out of range") + def _getslice(self, slice_object): + + l = [] + + for i in xrange(*slice_object.indices(len(self))): + l.append(self[i]) + return l + def __iter__(self): if self._clone_iter: return self._clone_iter() diff --git a/tests/test_iterator.py b/tests/test_iterator.py index 62ae75f..cd7732c 100644 --- a/tests/test_iterator.py +++ b/tests/test_iterator.py @@ -40,8 +40,6 @@ class TestFile(unittest.TestCase): def test_len(self): f = aaf.open(main_test_file) - - iterable = f.storage.mobs() assert len(iterable) == 199 @@ -55,10 +53,50 @@ def test_len(self): assert count == len(iterable) assert len(iterable) == 199 - def test_skip(self): + def test_negative_index(self): + f = aaf.open(main_test_file) + iterable = f.storage.mobs() + + + last_item = iterable[len(iterable)-1] + assert last_item == iterable[-1] + + try: + iterable[-10000] + except IndexError: + pass + else: + raise + + for i in xrange(len(iterable)): + assert iterable[i] == iterable[i-len(iterable)] + + def test_slice(self): f = aaf.open(main_test_file) + iterable = f.storage.mobs() + + + s = iterable[1:10] + l = [] + for i in xrange(1, 10): + l.append(iterable[i]) + + assert s == l + + s = iterable[100:-10:2] + l = [] + for i in xrange(100, len(iterable)-10, 2): + l.append(iterable[i]) + + assert s == l + + l = iterable[-1000: 1000] + + + def test_skip(self): + f = aaf.open(main_test_file) iterable = f.storage.mobs() assert len(iterable) == 199