From d9811ca548b66aa801fe90164b548149cf842822 Mon Sep 17 00:00:00 2001 From: Muthu Annamalai Date: Thu, 20 Nov 2025 12:09:03 -0800 Subject: [PATCH 1/2] [#106][testing]: update testin for async pickle db --- tests.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests.py b/tests.py index 798e69a..7405e87 100644 --- a/tests.py +++ b/tests.py @@ -1,3 +1,4 @@ +import pathlib import unittest import os import time @@ -9,12 +10,13 @@ # Adjust the import path if needed. For example, if 'pickledb' is your own module, # ensure the relative or absolute path matches your project structure. from pickledb import PickleDB, AsyncPickleDB - +from pathlib import Path +CURRENT_DIR = Path(__file__).absolute().parent class TestPickleDB(unittest.TestCase): def setUp(self): """Set up a PickleDB instance with a real file.""" - self.test_file = "test_pickledb.json" + self.test_file =CURRENT_DIR / "test_pickledb.json" self.db = PickleDB(self.test_file) def tearDown(self): @@ -123,7 +125,7 @@ def test_remove_non_string_key(self): class TestAsyncPickleDB(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): """Set up an AsyncPickleDB instance with a real file.""" - self.test_file = "test_async_pickledb.json" + self.test_file = CURRENT_DIR / "test_async_pickledb.json" if os.path.exists(self.test_file): os.remove(self.test_file) self.db = AsyncPickleDB(self.test_file) @@ -183,7 +185,9 @@ async def test_asave_and_reload(self): # Create a new AsyncPickleDB instance to verify persistence, # then use its inherited synchronous `get` method. new_db = AsyncPickleDB(self.test_file) - self.assertEqual(new_db.get("async_key"), "async_val") + await new_db.aload() + keyval = await new_db.aget("async_key") + self.assertEqual(keyval, "async_val") async def test_aset_non_string_key(self): """Test setting a non-string key asynchronously.""" From 8da39bda2d928332bb1e9d11c6658f770dec5fee Mon Sep 17 00:00:00 2001 From: Muthu Annamalai Date: Thu, 20 Nov 2025 12:38:58 -0800 Subject: [PATCH 2/2] [#105]: add dirty flag based save; improves save() operation to be conditional; e.g. stress testing becomes faster. Retrieved 1000000 key-value pairs in 0.22 seconds Dumped 1000000 key-value pairs to disk in 0.03 seconds --- pickledb.py | 41 +++++++++++++++++++++++++++++++++++------ tests.py | 10 +++++----- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/pickledb.py b/pickledb.py index ff6daca..7e0e6b5 100644 --- a/pickledb.py +++ b/pickledb.py @@ -33,13 +33,28 @@ import aiofiles import orjson +class ModifiedFlagImpl: + __slots__ = ('_dirty',) + def __init__(self): + self._dirty = False -class PickleDB: + def get_modified(self): + return self._dirty + + def set_modified(self): + self._dirty = True + return self._dirty + + def clear_modified(self): + self._dirty = False + return self._dirty + +class PickleDB(ModifiedFlagImpl): """ A barebones orjson-based key-value store with essential methods: set, get, save, remove, purge, and all. """ - + __slots__ = ('location','db') def __init__(self, location): """ Initialize the PickleDB object. @@ -47,14 +62,17 @@ def __init__(self, location): Args: location (str): Path to the JSON file. """ + super(PickleDB,self).__init__() self.location = os.path.expanduser(location) - self._load() + self._load() #defines .db + def __setitem__(self, key, value): """ Wraps the `set` method to allow `db[key] = value`. See `set` method for details. """ + self.set_modified() return self.set(key, value) def __getitem__(self, key): @@ -94,6 +112,8 @@ def _load(self): raise RuntimeError(f"{e}\nFailed to load database.") else: self.db = {} + self.clear_modified() + return def save(self, option=0): """ @@ -116,6 +136,7 @@ def save(self, option=0): with open(temp_location, "wb") as temp_file: temp_file.write(orjson.dumps(self.db, option=option)) os.replace(temp_location, self.location) + self.clear_modified() return True except Exception as e: print(f"Failed to save database: {e}") @@ -140,6 +161,7 @@ def set(self, key, value): """ key = str(key) if not isinstance(key, str) else key self.db[key] = value + self.set_modified() return True def remove(self, key): @@ -157,6 +179,7 @@ def remove(self, key): key = str(key) if not isinstance(key, str) else key if key in self.db: del self.db[key] + self.set_modified() return True return False @@ -168,6 +191,7 @@ def purge(self): bool: True if the operation succeeds. """ self.db.clear() + self.set_modified() return True def get(self, key): @@ -195,13 +219,14 @@ def all(self): return list(self.db.keys()) -class AsyncPickleDB: +class AsyncPickleDB(ModifiedFlagImpl): """ A fully asynchronous orjson-based key-value store. Provides async load, save, and CRUD operations with file locking. """ - + __slots__ = ('_lock','location', 'db') def __init__(self, location): + super(AsyncPickleDB,self).__init__() self.location = os.path.expanduser(location) self._lock = asyncio.Lock() self.db = {} @@ -224,6 +249,7 @@ async def aload(self): async with aiofiles.open(self.location, "rb") as f: content = await f.read() self.db = orjson.loads(content) + self.clear_modified() except Exception as e: raise RuntimeError(f"{e}\nFailed to load database.") else: @@ -239,6 +265,7 @@ async def asave(self, option=0): async with aiofiles.open(temp_location, "wb") as temp_file: await temp_file.write(orjson.dumps(self.db, option=option)) await asyncio.to_thread(os.replace, temp_location, self.location) + self.clear_modified() return True except Exception as e: print(f"Failed to save database: {e}") @@ -247,6 +274,7 @@ async def asave(self, option=0): async def aset(self, key, value): async with self._lock: self.db[str(key)] = value + self.set_modified() return True async def aget(self, key): @@ -255,6 +283,7 @@ async def aget(self, key): async def aremove(self, key): async with self._lock: + self.set_modified() return self.db.pop(str(key), None) is not None async def aall(self): @@ -263,6 +292,6 @@ async def aall(self): async def apurge(self): async with self._lock: + self.set_modified() self.db.clear() return True - diff --git a/tests.py b/tests.py index 7405e87..58cac19 100644 --- a/tests.py +++ b/tests.py @@ -1,17 +1,14 @@ -import pathlib import unittest import os import time import signal import asyncio -import aiofiles -import orjson # Adjust the import path if needed. For example, if 'pickledb' is your own module, # ensure the relative or absolute path matches your project structure. -from pickledb import PickleDB, AsyncPickleDB from pathlib import Path -CURRENT_DIR = Path(__file__).absolute().parent +from pickledb import PickleDB, AsyncPickleDB +CURRENT_DIR = Path(__file__).parent.absolute() class TestPickleDB(unittest.TestCase): def setUp(self): @@ -56,7 +53,9 @@ def test_stress_operation(self): # Measure dump performance start_time = time.time() + self.assertTrue( self.db.get_modified() ) self.db.save() + self.assertFalse( self.db.get_modified() ) dump_time = time.time() - start_time print(f"Dumped {num_docs} key-value pairs to disk in {dump_time:.2f} seconds") @@ -99,6 +98,7 @@ def test_all_keys(self): def test_dump_and_reload(self): """Test dumping the database to disk and reloading it.""" self.db.set("key1", "value1") + self.assertTrue(self.db.get_modified()) self.db.save() reloaded_db = PickleDB(self.test_file) self.assertEqual(reloaded_db.get("key1"), "value1")