diff --git a/ordered_set/__init__.py b/ordered_set/__init__.py index ccd1cbf..e7bb09a 100644 --- a/ordered_set/__init__.py +++ b/ordered_set/__init__.py @@ -266,6 +266,9 @@ def pop(self, index: int = -1) -> T: elem = self.items[index] del self.items[index] del self.map[elem] + # Rebuild map so indices stay consistent for remaining items. + self.map = {item: idx for (idx, item) in enumerate(self.items)} + return elem def discard(self, key: T) -> None: diff --git a/test/test_ordered_set.py b/test/test_ordered_set.py index 6efe0a9..1b36759 100644 --- a/test/test_ordered_set.py +++ b/test/test_ordered_set.py @@ -162,6 +162,28 @@ def test_pop(): pytest.raises(KeyError, set1.pop) +def test_pop_with_index(): + set1 = OrderedSet("abcde") + elem = set1.pop(0) + + assert elem == "a" + assert list(set1) == ["b", "c", "d", "e"] + assert set1.index("b") == 0 + assert set1[1] == "c" + + elem = set1.pop(1) + + assert elem == "c" + assert list(set1) == ["b", "d", "e"] + assert set1.index("d") == 1 + assert set1[1] == "d" + + elem = set1.pop() + + assert elem == "e" + assert list(set1) == ["b", "d"] + + def test_getitem_type_error(): set1 = OrderedSet("ab") with pytest.raises(TypeError):