Merge pull request #76 from kaznak/202107.kaznak.closed-order-handling

Add a method for deletion to DataStore (#75)
pybotters · Aug 15, 2021 · f25934c · f25934c
2 parents a8d2e43 + dd6c3fc
commit f25934c
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 0 deletions.
diff --git a/pybotters/store.py b/pybotters/store.py
@@ -133,6 +133,20 @@ def get(self, item: Item) -> Optional[Item]:
                 if keyhash in self._index:
                     return self._data[self._index[keyhash]]
 
+    def _pop(self, item: Item) -> Optional[Item]:
+        if self._keys:
+            try:
+                keyitem = {k: item[k] for k in self._keys}
+            except KeyError:
+                pass
+            else:
+                keyhash = self._hash(keyitem)
+                if keyhash in self._index:
+                    ret = self._data[self._index[keyhash]]
+                    del self._data[self._index[keyhash]]
+                    del self._index[keyhash]
+                    return ret
+
     def find(self, query: Item = {}) -> List[Item]:
         if query:
             return [
@@ -143,6 +157,20 @@ def find(self, query: Item = {}) -> List[Item]:
         else:
             return list(self)
 
+    def _find_and_delete(self, query: Item = {}) -> List[Item]:
+        if query:
+            ret = [
+                item
+                for item in self
+                if all(k in item and query[k] == item[k] for k in query)
+            ]
+            self._delete(ret)
+            return ret
+        else:
+            ret = list(self)
+            self._clear()
+            return ret
+
     def _set(self) -> None:
         for event in self._events:
             event.set()

diff --git a/tests/test_store.py b/tests/test_store.py
@@ -161,6 +161,53 @@ def test_delete():
     assert len(ds3._index) == 1000
 
 
+def test_pop():
+    data = [{'foo': f'bar{i}'} for i in range(1000)]
+
+    ds1 = pybotters.store.DataStore(keys=['foo'], data=data)
+    assert ds1._pop({'foo': 'bar500'}) == {'foo': 'bar500'}
+    assert ds1.get({'foo': 'bar500'}) is None
+    assert ds1._pop({'foo': 'bar9999'}) is None
+
+    ds2 = pybotters.store.DataStore(data=data)
+    assert ds2._pop({'foo': 'bar500'}) is None
+
+
+def test_find_and_delete():
+    data = [{'foo': f'bar{i}', 'mod': i % 2} for i in range(1000)]
+    query = {'mod': 1}
+    invalid = {'mod': -1}
+
+    ds1 = pybotters.store.DataStore(keys=['foo'], data=data)
+    ret1 = ds1._find_and_delete()
+    # return value
+    assert isinstance(ret1, list)
+    assert len(ret1) == 1000
+    # data store
+    assert len(ds1._data) == 0
+    assert len(ds1._index) == 0
+
+    ds2 = pybotters.store.DataStore(keys=['foo'], data=data)
+    ret2 = ds2._find_and_delete(query)
+    # return value
+    assert isinstance(ret2, list)
+    assert len(ret2) == 500
+    assert all(map(lambda record: 1 == record['mod'], ret2))
+    # data store
+    assert len(ds2._data) == 500
+    assert all(map(lambda record: 0 == record['mod'], ds2._data.values()))
+    assert len(ds2._index) == 500
+
+    ds3 = pybotters.store.DataStore(keys=['foo'], data=data)
+    ret3 = ds3._find_and_delete(invalid)
+    # return value
+    assert isinstance(ret3, list)
+    assert len(ret3) == 0
+    # data store
+    assert len(ds3._data) == 1000
+    assert len(ds3._index) == 1000
+
+
 def test_clear():
     data = [{'foo': f'bar{i}'} for i in range(1000)]
     ds = pybotters.store.DataStore(keys=['foo'], data=data)