diff --git a/.gitignore b/.gitignore index 249963d..67d5de1 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ *.iml -.idea/ \ No newline at end of file +.idea/ +*.db \ No newline at end of file diff --git a/hexlib/db.py b/hexlib/db.py index e69de29..86c602f 100644 --- a/hexlib/db.py +++ b/hexlib/db.py @@ -0,0 +1,82 @@ +import sqlite3 + + +class PersistentState: + """Quick and dirty persistent dict-like SQLite wrapper""" + + def __init__(self, dbfile="state.db"): + self.dbfile = dbfile + + def __getitem__(self, table): + return Table(self, table) + + +class Table: + def __init__(self, state, table): + self._state = state + self._table = table + + def __iter__(self): + with sqlite3.connect(self._state.dbfile) as conn: + conn.row_factory = sqlite3.Row + try: + cur = conn.execute("SELECT * FROM %s" % (self._table,)) + for row in cur: + yield dict(row) + except: + return None + + def __getitem__(self, item): + with sqlite3.connect(self._state.dbfile) as conn: + conn.row_factory = sqlite3.Row + try: + cur = conn.execute("SELECT * FROM %s WHERE id=?" % (self._table,), (item,)) + + row = cur.fetchone() + if row: + return dict(row) + except: + return None + + def __setitem__(self, key, value): + + with sqlite3.connect(self._state.dbfile) as conn: + conn.row_factory = sqlite3.Row + + sql = "INSERT INTO %s (id,%s) VALUES ('%s',%s)" % \ + (self._table, ",".join(value.keys()), key, ",".join(_sqlite_value(v) for v in value.values())) + + try: + conn.execute(sql) + 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) + + except sqlite3.IntegrityError: + conn.execute( + "UPDATE %s SET (%s) = (%s) WHERE id=?" % + (self._table, ",".join(value.keys()), ",".join(_sqlite_value(v) for v in value.values())), + (key,) + ) + + +def _sqlite_type(value): + if isinstance(value, int): + return "integer" + if isinstance(value, float): + return "real" + if isinstance(value, bytes): + return "blob" + return "text" + + +def _sqlite_value(value): + if isinstance(value, str): + return '"%s"' % value + if isinstance(value, bytes): + return _sqlite_value(value.decode("unicode_escape")) + return str(value) + diff --git a/setup.py b/setup.py index 982402a..77252d8 100644 --- a/setup.py +++ b/setup.py @@ -2,12 +2,12 @@ from setuptools import setup setup( name="hexlib", - version="1.0", + version="1.1", description="Misc utility methods", author="simon987", author_email="me@simon987.net", packages=["hexlib"], install_requires=[ - "ImageHash", "influxdb", + "ImageHash", "influxdb", "siphash" ] )