From 23f23f45bec173955b3f192e7cb25362ba43e4f5 Mon Sep 17 00:00:00 2001 From: David Chiu Date: Thu, 6 Jun 2024 20:58:55 +0800 Subject: [PATCH] chore: fix closing of LMDB cursor in LmdbStore --- sync_crawler/store/lmdb_store.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/sync_crawler/store/lmdb_store.py b/sync_crawler/store/lmdb_store.py index 9413fa0..5f32cfa 100644 --- a/sync_crawler/store/lmdb_store.py +++ b/sync_crawler/store/lmdb_store.py @@ -3,6 +3,7 @@ import pickle import uuid from collections.abc import Callable, Iterable +from contextlib import closing from typing import override import lmdb @@ -40,7 +41,7 @@ def put(self, news: Iterable[News]): with ( self._env.begin(write=True) as txn, - txn.cursor() as cur, + closing(txn.cursor()) as cur, ): cur.putmulti(key_value_pairs) @@ -48,8 +49,11 @@ def put(self, news: Iterable[News]): def pop(self, nums=1) -> Iterable[News]: values: Iterable[News] = [] - with self._env.begin(write=True) as txn: - for key, value in itertools.islice(txn.cursor(), nums): + with ( + self._env.begin(write=True) as txn, + closing(txn.cursor()) as cur, + ): + for key, value in itertools.islice(cur, nums): values.append(pickle.loads(value)) txn.delete(key)