AYMENJD · chenray844 · Aug 10, 2023 · Aug 11, 2023
diff --git a/rocksdb/client.py b/rocksdb/client.py
@@ -5,7 +5,8 @@
 from concurrent.futures import ThreadPoolExecutor
 from .rocksdb_ext import RocksDBext, Response
 
-import rocksdb, asyncio
+import rocksdb
+import asyncio
 
 logger = getLogger(__name__)
 
@@ -143,7 +144,7 @@ async def get(self, options: ReadOptions, key: str) -> Response:
 
         return await future
 
-    async def put(self, options: WriteOptions, key: str, value: str) -> Response:
+    async def put(self, options: WriteOptions, key: str, value: str | bytes) -> Response:
         """Set the database entry for `key` to `value`
 
         Args:
@@ -166,8 +167,8 @@ async def put(self, options: WriteOptions, key: str, value: str) -> Response:
 
         if not isinstance(key, str):
             raise TypeError("key must be str")
-        elif not isinstance(value, str):
-            raise TypeError("value must be str")
+        elif not isinstance(value, (str, bytes)):
+            raise TypeError("value must be str | bytes")
         elif not isinstance(options, WriteOptions):
             raise TypeError(f"Invalid class '{type(options).__name__}'")
 
@@ -276,7 +277,8 @@ async def getOptions(self) -> Response:
             `Response`
         """
 
-        future = self.loop.run_in_executor(self.executer, self.__rocksdb.GetOptions)
+        future = self.loop.run_in_executor(
+            self.executer, self.__rocksdb.GetOptions)
 
         return await future
 
@@ -369,7 +371,8 @@ async def flush(self, options: FlushOptions) -> Response:
         if not isinstance(options, FlushOptions):
             raise TypeError(f"Invalid class '{type(options).__name__}'")
 
-        future = self.loop.run_in_executor(self.executer, self.__rocksdb.Flush, options)
+        future = self.loop.run_in_executor(
+            self.executer, self.__rocksdb.Flush, options)
 
         return await future
 

diff --git a/rocksdb/ext/extension.cpp b/rocksdb/ext/extension.cpp
@@ -28,7 +28,8 @@ PYBIND11_MODULE(rocksdb_ext, m) {
         .def(py::init<rocksdb::Status &, std::string *>())
         .def_readwrite("status", &Response::status)
         .def_readwrite("options", &Response::options)
-        .def_readwrite("value", &Response::value);
+        .def_readwrite("value", &Response::value)
+        .def("value_bytes", &Response::value_bytes);
 
     // RocksDB options aka rocksdb::Options
     py::class_<rocksdb::Options>(m, "_Options")

diff --git a/rocksdb/ext/rocksdb.hpp b/rocksdb/ext/rocksdb.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <exception>
+#include <iostream>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <rocksdb/db.h>
@@ -24,6 +25,10 @@ class Response {
 
         this->status = s;
     }
+
+    const py::bytes value_bytes() {
+        return py::bytes(this->value);
+    };
 };
 
 class RocksDB {