diff --git a/CHANGELOG.md b/CHANGELOG.md index d42561f..a14bfda 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ Versioning]. - Improve performance by reducing the AES plaintext size. Note: this changes the ordering of integers for existing seeds, which may be a breaking change depending on the use case. +- Add support for negative indexes. ## [1.1.0] - 2022-12-29 diff --git a/src/shuffled/core.py b/src/shuffled/core.py index 810e44f..cd9b504 100644 --- a/src/shuffled/core.py +++ b/src/shuffled/core.py @@ -46,6 +46,8 @@ def __getitem__(self, index: slice) -> Sequence[int]: ... def __getitem__(self, index: Union[int, slice]) -> Union[int, Sequence[int]]: if isinstance(index, slice): raise NotImplementedError + if index < 0: + index += self._size if index < 0 or index >= self._size: raise IndexError return self._encryptor.encrypt(index) diff --git a/tests/test_shuffled.py b/tests/test_shuffled.py index bf91ddf..c763f23 100644 --- a/tests/test_shuffled.py +++ b/tests/test_shuffled.py @@ -31,3 +31,15 @@ def test_normal(self, range_size): ) def test_seed(self, seed, expected): assert list(Shuffled(20, seed)) == expected + + def test_negative_index(self): + s = Shuffled(10, seed=b"\x00") + + assert s[-1] == s[9] + assert s[-10] == s[0] + + def test_negative_index_out_of_bounds(self): + s = Shuffled(10, seed=b"\x00") + + with pytest.raises(IndexError): + s[-11]