From dd676e2f73699a72c04c51e77f55a950bd10b801 Mon Sep 17 00:00:00 2001 From: Robert Asztalos Date: Tue, 31 May 2022 13:18:34 +0300 Subject: [PATCH] Rebuild self.map after .pop() --- ordered_set/__init__.py | 4 +++- test/test_ordered_set.py | 15 +++++++++++---- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/ordered_set/__init__.py b/ordered_set/__init__.py index e86c70e..96a9e45 100644 --- a/ordered_set/__init__.py +++ b/ordered_set/__init__.py @@ -252,7 +252,8 @@ def index(self, key): def pop(self, index=-1) -> T: """ - Remove and return item at index (default last). + Remove and return item at index (default last), updating + self.map accordingly. Raises KeyError if the set is empty. Raises IndexError if index is out of range. @@ -268,6 +269,7 @@ def pop(self, index=-1) -> T: elem = self.items[index] del self.items[index] del self.map[elem] + 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..6adbc26 100644 --- a/test/test_ordered_set.py +++ b/test/test_ordered_set.py @@ -151,13 +151,20 @@ def test_update(): def test_pop(): - set1 = OrderedSet("ab") - elem = set1.pop() + set1 = OrderedSet("abc") + elem = set1.pop(0) + + assert elem == "a" + assert set1.index('b') == 0 + assert set1.index('c') == 1 - assert elem == "b" elem = set1.pop() + assert elem == "c" + assert set1.index('b') == 0 - assert elem == "a" + elem = set1.pop() + assert elem == "b" + assert not set1.map pytest.raises(KeyError, set1.pop)