From 86e07900ea3e7c01ac3b5207e7c7c9a73b001bba Mon Sep 17 00:00:00 2001
From: simon987 <me@simon987.net>
Date: Sun, 29 Dec 2019 17:13:35 -0500
Subject: [PATCH] Handle raw bytes

---
 hexlib/db.py | 33 +++++++++++++++++++++++----------
 setup.py     |  2 +-
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/hexlib/db.py b/hexlib/db.py
index 3284375..6384330 100644
--- a/hexlib/db.py
+++ b/hexlib/db.py
@@ -1,11 +1,15 @@
+import base64
 import sqlite3
 
 
 class PersistentState:
     """Quick and dirty persistent dict-like SQLite wrapper"""
 
-    def __init__(self, dbfile="state.db"):
+    def __init__(self, dbfile="state.db", **dbargs):
         self.dbfile = dbfile
+        if dbargs is None:
+            dbargs = {"timeout": 30000}
+        self.dbargs = dbargs
 
     def __getitem__(self, table):
         return Table(self, table)
@@ -17,7 +21,7 @@ class Table:
         self._table = table
 
     def __iter__(self):
-        with sqlite3.connect(self._state.dbfile) as conn:
+        with sqlite3.connect(self._state.dbfile, **self._state.dbargs) as conn:
             conn.row_factory = sqlite3.Row
             try:
                 cur = conn.execute("SELECT * FROM %s" % (self._table,))
@@ -27,38 +31,42 @@ class Table:
                 return None
 
     def __getitem__(self, item):
-        with sqlite3.connect(self._state.dbfile) as conn:
+        with sqlite3.connect(self._state.dbfile, **self._state.dbargs) as conn:
             conn.row_factory = sqlite3.Row
             try:
+                col_types = conn.execute("PRAGMA table_info(%s)" % self._table).fetchall()
                 cur = conn.execute("SELECT * FROM %s WHERE id=?" % (self._table,), (item,))
 
                 row = cur.fetchone()
                 if row:
-                    return dict(row)
+                    return dict(
+                        (col[0], _deserialize(row[col[0]], col_types[i]["type"]))
+                        for i, col in enumerate(cur.description)
+                    )
             except:
                 return None
 
     def __setitem__(self, key, value):
 
-        with sqlite3.connect(self._state.dbfile) as conn:
+        with sqlite3.connect(self._state.dbfile, **self._state.dbargs) as conn:
             conn.row_factory = sqlite3.Row
 
             sql = "INSERT INTO %s (id,%s) VALUES ('%s',%s)" % \
                   (self._table, ",".join(value.keys()), key, ",".join("?" for _ in value.values()))
             try:
-                conn.execute(sql, list(_sqlite_value(v) for v in value.values()))
+                conn.execute(sql, list(_serialize(v) for v in value.values()))
             except sqlite3.OperationalError:
                 conn.execute(
                     "create table if not exists %s (id text primary key,%s)" %
                     (self._table, ",".join("%s %s" % (k, _sqlite_type(v)) for k, v in value.items()))
                 )
-                conn.execute(sql, list(_sqlite_value(v) for v in value.values()))
+                conn.execute(sql, list(_serialize(v) for v in value.values()))
 
             except sqlite3.IntegrityError:
                 sql = "UPDATE %s SET (%s) = (%s) WHERE id=?" \
                       % (self._table, ",".join(value.keys()), ",".join("?" for _ in value.values()))
                 args = [key]
-                args.extend(_sqlite_value(v) for v in value.values())
+                args.extend(_serialize(v) for v in value.values())
                 conn.execute(
                     sql,
                     args
@@ -75,8 +83,13 @@ def _sqlite_type(value):
     return "text"
 
 
-def _sqlite_value(value):
+def _serialize(value):
     if isinstance(value, bytes):
-        return _sqlite_value(value.decode("unicode_escape"))
+        return base64.b64encode(value)
     return str(value)
 
+
+def _deserialize(value, col_type):
+    if col_type == "blob":
+        return base64.b64decode(value)
+    return value
diff --git a/setup.py b/setup.py
index 7e8cc46..1acc8eb 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup
 
 setup(
     name="hexlib",
-    version="1.5",
+    version="1.6",
     description="Misc utility methods",
     author="simon987",
     author_email="me@simon987.net",