diff --git a/statemachine/orderedset.py b/statemachine/orderedset.py index 71d49779..29f6fced 100644 --- a/statemachine/orderedset.py +++ b/statemachine/orderedset.py @@ -62,10 +62,14 @@ class OrderedSet(MutableSet[T]): >>> s[2] 3 - >>> eval(repr(OrderedSet(['a', 'b', 'c']))) - OrderedSet(['a', 'b', 'c']) + >>> s[-1] + 3 + >>> s[-3] + 1 + >>> eval(repr(OrderedSet(['a', 'b', 'c']))) + OrderedSet(['a', 'b', 'c']) """ @@ -84,10 +88,15 @@ def discard(self, x: T) -> None: self._d.pop(x, None) def __getitem__(self, index) -> T: + original_index = index + if index < 0: + index += len(self) + if index < 0: + raise IndexError(f"index {original_index} out of range") try: return next(itertools.islice(self._d, index, index + 1)) except StopIteration as e: - raise IndexError(f"index {index} out of range") from e + raise IndexError(f"index {original_index} out of range") from e def __contains__(self, x: object) -> bool: return self._d.__contains__(x) diff --git a/tests/test_state.py b/tests/test_state.py index c2f97c67..2e2d7f1c 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -56,6 +56,14 @@ def test_ordered_set_getitem(): assert s[2] == 30 +def test_ordered_set_getitem_negative_index(): + """OrderedSet supports negative index access like Python sequences.""" + s = OrderedSet([10, 20, 30]) + assert s[-1] == 30 + assert s[-2] == 20 + assert s[-3] == 10 + + def test_ordered_set_getitem_out_of_range(): """OrderedSet raises IndexError for out-of-range index.""" s = OrderedSet([10, 20]) @@ -63,6 +71,13 @@ def test_ordered_set_getitem_out_of_range(): s[5] +def test_ordered_set_getitem_negative_out_of_range(): + """OrderedSet raises IndexError for negative out-of-range index.""" + s = OrderedSet([10, 20]) + with pytest.raises(IndexError, match=r"index -3 out of range"): + s[-3] + + def test_ordered_set_union(): """OrderedSet.union returns new set with elements from both.""" s1 = OrderedSet([1, 2])