Merge pull request #4 from labteral/develop

v2

README.md

@@ -0,0 +1,32 @@
+# easyrocks
+A [`python-rocksdb`](https://github.com/twmht/python-rocksdb) wrapper for a more comfortable interaction with RocksDB.
+
+> Keys must be of type `bytes`. Values can be nested structures or complex objects. Value serialization is automatically performed with MessagePack if possible, otherwise with Pickle.
+
+## Usage
+```python
+from easyrocks import RocksDB
+
+db = RocksDB(path='./rocksdb', read_only=False)
+
+key = b'key1'
+db.put(key, 'one')
+db.get(key)
+db.exists(key)
+
+for key, value in db.scan(prefix=None, start_key=None, stop_key=None, reversed_scan=False):
+  print(key, value)
+```
+
+## Utils
+There are some useful functions to transform a key to and from `bytes`:
+```python
+from easyrocks.utils import (
+  str_to_bytes
+  bytes_to_str
+  int_to_bytes
+  bytes_to_int
+  str_to_padded_bytes
+  int_to_padded_bytes
+)
+```

easyrocks/__init__.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+from rocksdb import *
 from .easyrocks import *
 
-__version__ = '0.0.17a'
+__version__ = '2.214.0'

easyrocks/easyrocks.py

@@ -3,13 +3,18 @@
 
 import gc
 from . import utils
-from rocksdb import *
-
-ALLOWED_KEY_TYPES = (int, str)
+from rocksdb import DB, Options, WriteBatch, BackupEngine
+from typing import Dict, Generator
 
 
 class RocksDB:
-    def __init__(self, path: str = './rocksdb', opts: dict = None, read_only: bool = False):
+
+    ALLOWED_KEY_TYPES = set([bytes, type(None)])
+
+    def __init__(self,
+                 path: str = './rocksdb',
+                 opts: Dict = None,
+                 read_only: bool = False):
         self._path = path
 
         if opts is None:
@@ -24,7 +29,7 @@ def path(self) -> str:
         return self._path
 
     @property
-    def opts(self) -> dict:
+    def opts(self) -> Dict:
         return dict(self._opts)
 
     @property
@@ -35,7 +40,7 @@ def read_only(self) -> bool:
     def db(self) -> DB:
         return self._db
 
-    def reload(self, opts: dict = None, read_only: bool = None):
+    def reload(self, opts: Dict = None, read_only: bool = None):
         if opts is None:
             opts = self._opts
 
@@ -51,70 +56,82 @@ def reload(self, opts: dict = None, read_only: bool = None):
 
         self._db = DB(self._path, rocks_opts, read_only=read_only)
 
-    def put(self, key, value, write_batch=None):
-        if not isinstance(key, ALLOWED_KEY_TYPES):
+    def put(self, key: bytes, value, write_batch=None):
+        if type(key) not in self.ALLOWED_KEY_TYPES:
             raise TypeError
 
         if value is None:
             raise ValueError
 
-        key_bytes = utils._get_key_bytes(key)
-        value_bytes = utils.to_bytes(value)
+        value_bytes = utils.pack(value)
 
         if write_batch is not None:
-            write_batch.put(key_bytes, value_bytes)
+            write_batch.put(key, value_bytes)
         else:
-            self._db.put(key_bytes, value_bytes, sync=True)
+            self._db.put(key, value_bytes, sync=True)
+
+    def get(self, key: bytes):
+        if type(key) not in self.ALLOWED_KEY_TYPES:
+            raise TypeError
 
-    def get(self, key):
-        key_bytes = utils._get_key_bytes(key)
-        value_bytes = self._db.get(key_bytes)
+        value_bytes = self._db.get(key)
 
         if value_bytes is not None:
-            return utils.to_object(value_bytes)
+            return utils.unpack(value_bytes)
 
-    def exists(self, key):
+    def exists(self, key: bytes) -> bool:
         if self.get(key) is not None:
             return True
         return False
 
-    def delete(self, key):
-        key_bytes = utils.str_to_bytes(key)
-        self._db.delete(key_bytes, sync=True)
+    def delete(self, key: bytes, write_batch: WriteBatch = None):
+        if type(key) not in self.ALLOWED_KEY_TYPES:
+            raise TypeError
 
-    def commit(self, write_batch):
         if write_batch is not None:
-            self._db.write(write_batch, sync=True)
+            write_batch.delete(key)
+        else:
+            self._db.delete(key, sync=True)
 
-    def scan(self, prefix=None, start_key=None, end_key=None, reversed_scan=False):
-        iterator = self._db.iterkeys()
+    def commit(self, write_batch: WriteBatch):
+        if write_batch is None:
+            raise ValueError
+        self._db.write(write_batch, sync=True)
+
+    def scan(self,
+             prefix: bytes = None,
+             start_key: bytes = None,
+             stop_key: bytes = None,
+             reversed_scan: bool = False) -> Generator:
 
+        for key in [prefix, start_key, stop_key]:
+            if type(key) not in self.ALLOWED_KEY_TYPES:
+                raise TypeError
+
+        iterator = self._db.iterkeys()
         if prefix is None and start_key is None:
             if reversed_scan:
                 iterator.seek_to_last()
             else:
                 iterator.seek_to_first()
+        elif prefix is not None:
+            iterator.seek(prefix)
         else:
-            if prefix is not None:
-                prefix_bytes = utils.str_to_bytes(prefix)
-            else:
-                prefix_bytes = utils.str_to_bytes(start_key)
-            iterator.seek(prefix_bytes)
+            iterator.seek(start_key)
 
         if reversed_scan:
             iterator = reversed(iterator)
 
-        for key_bytes in iterator:
-            key = utils.bytes_to_str(key_bytes)
-
+        for key in iterator:
             if prefix is not None and key[:len(prefix)] != prefix:
                 return
 
-            if end_key is not None and key > end_key:
+            if stop_key is not None and key > stop_key:
                 return
 
-            value_bytes = self._db.get(key_bytes)
-            value = utils.to_object(value_bytes)
+            value_bytes = self._db.get(key)
+            value = utils.unpack(value_bytes)
+
             yield key, value
 
     def close(self):
@@ -142,4 +159,4 @@ def restore_backup(self, path: str, backup_id: int = None):
             backup_engine.restore_latest_backup(self._path, self._path)
         else:
             backup_engine.restore_backup(backup_id, self._path, self._path)
-        self.reload()
+        self.reload()

easyrocks/utils.py

@@ -5,59 +5,50 @@
 import msgpack
 
 
-def to_bytes(value) -> bytes:
+def pack(value) -> bytes:
     try:
         return msgpack.packb(value)
     except TypeError:
         return pickle.dumps(value, protocol=5)
 
 
-def to_object(value: bytes):
+def unpack(value: bytes):
     try:
         return msgpack.unpackb(value)
     except Exception:
         return pickle.loads(value)
 
 
-def str_to_bytes(string):
-    if string is None:
-        return
-    return bytes(string, 'utf-8')
+def str_to_bytes(value: str) -> bytes:
+    return bytes(value, 'utf-8')
 
 
-def bytes_to_str(bytes_string):
-    if bytes_string is None:
-        return
-    return bytes_string.decode('utf-8')
+def bytes_to_str(value: bytes) -> str:
+    if value is not None:
+        return value.decode('utf-8')
 
 
-def get_padded_int(integer, size=32, left=False, right=False):
-    integer_string = str(integer)
-    return get_padded_str(integer_string, size, left, right)
+def _big_endian_to_int(value: bytes) -> int:
+    return int.from_bytes(value, "big")
 
 
-def get_padded_str(key, size=64, left=False, right=False):
-    if not left and not right:
-        left = True
-    zeros = size - len(key)
-    if zeros < 0:
-        raise ValueError
-    if left:
-        new_key = f"{zeros * '0'}{key}"
-    else:
-        new_key = f"{key}{zeros * '0'}"
-    return new_key
+def _int_to_big_endian(value: int) -> bytes:
+    return value.to_bytes((value.bit_length() + 7) // 8 or 1, 'big')
 
 
-def int_to_bytes(integer):
-    return str_to_bytes(get_padded_int(integer))
+def int_to_bytes(value: int):
+    return _int_to_big_endian(value)
 
 
-def _get_key_bytes(key):
-    if isinstance(key, int):
-        key_bytes = int_to_bytes(key)
-    elif isinstance(key, str):
-        key_bytes = str_to_bytes(key)
-    else:
-        raise TypeError
-    return key_bytes
+def bytes_to_int(value: bytes):
+    return _big_endian_to_int(value)
+
+
+def str_to_padded_bytes(value: str, length: int) -> bytes:
+    str_bytes = str_to_bytes(value)
+    return str_bytes.rjust(length, b'\x00')
+
+
+def int_to_padded_bytes(value: int, lenght: int) -> bytes:
+    int_bytes = int_to_bytes(value)
+    return int_bytes.rjust(lenght, b'\x00')

tests/main.py

@@ -3,6 +3,7 @@
 
 import easyrocks
 from easyrocks import RocksDB
+from easyrocks.utils import str_to_bytes
 
 print(f'easyrocks v{easyrocks.__version__}')
 db = RocksDB(path='/tmp/rocksdb')
@@ -22,9 +23,10 @@
 
 # PUT
 for i in range(6):
-    db.put(f'key{i + 1}', f'value{i + 1}')
-    assert db.get(f'key{i + 1}') == f'value{i + 1}'
-    assert db.exists(f'key{i + 1}')
+    key = str_to_bytes(f'key{i + 1}')
+    db.put(key, f'value{i + 1}')
+    assert db.get(key) == f'value{i + 1}'
+    assert db.exists(key)
 
 # CREATE TWO BACKUPS
 db.backup(path=backup_path)
@@ -45,9 +47,10 @@
 assert number_of_backup - 1 == len(backup_info)
 
 # DELETE
-db.delete('key6')
-assert not db.get('key6')
-assert not db.exists('key6')
+key = str_to_bytes('key6')
+db.delete(key)
+assert not db.get(key)
+assert not db.exists(key)
 
 # REGULAR SCAN
 index = 1
@@ -65,27 +68,34 @@
 
 # START KEY
 index = 2
-for _, value in db.scan(start_key='key2'):
+start_key = str_to_bytes('key2')
+for _, value in db.scan(start_key=start_key):
     assert value == f'value{index}'
     index += 1
 assert index == 6
 
 # START & STOP KEYS
 index = 2
-for _, value in db.scan(start_key='key2', end_key='key3'):
+start_key = str_to_bytes('key2')
+stop_key = str_to_bytes('key3')
+for _, value in db.scan(start_key=start_key, stop_key=stop_key):
     assert value == f'value{index}'
     index += 1
 assert index == 4
 
 # START & STOP KEYS ARE THE SAME
-for _, value in db.scan(start_key='key3', end_key='key3'):
+start_key = str_to_bytes('key3')
+stop_key = str_to_bytes('key3')
+for _, value in db.scan(start_key=start_key, stop_key=stop_key):
     assert value == 'value3'
 
 # START & STOP KEYS DO NOT EXIST
 index = 1
-for _, value in db.scan(start_key='key0', end_key='key9'):
+start_key = str_to_bytes('key0')
+stop_key = str_to_bytes('key9')
+for _, value in db.scan(start_key=start_key, stop_key=stop_key):
     assert value == f'value{index}'
     index += 1
 assert index == 6
 
-print('OK!')
+print('OK!')