Spaces:
Running
Running
File size: 1,170 Bytes
fa9fcab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
from pathlib import Path
from typing import Any, Iterator, Optional
# rocksdict is a lightweight C wrapper around RocksDB for Python, pylint may not recognize it
# pylint: disable=no-name-in-module
from rocksdict import Rdict
class RocksDBCache:
def __init__(self, cache_dir: str):
self.db_path = Path(cache_dir)
self.db = Rdict(str(self.db_path))
def get(self, key: str) -> Optional[Any]:
return self.db.get(key)
def set(self, key: str, value: Any):
self.db[key] = value
def delete(self, key: str):
try:
del self.db[key]
except KeyError:
# If the key does not exist, do nothing (deletion is idempotent for caches)
pass
def close(self):
if hasattr(self, "db") and self.db is not None:
self.db.close()
self.db = None
def __del__(self):
# Ensure the database is closed when the object is destroyed
self.close()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
def __iter__(self) -> Iterator[str]:
return iter(self.db.keys())
|