diff --git a/pybotters/store.py b/pybotters/store.py index e4eebb3..a3081e2 100644 --- 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 index f20e262..ec3f5ab 100644 --- 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)