mirror of
https://github.com/moparisthebest/SickRage
synced 2024-12-15 04:22:17 -05:00
162 lines
3.6 KiB
Python
162 lines
3.6 KiB
Python
|
"""
|
||
|
by default, things are stored in a "shelf" table
|
||
|
|
||
|
>>> d = SQLiteShelf("test.sdb")
|
||
|
|
||
|
you can put multiple shelves into a single SQLite database
|
||
|
|
||
|
>>> e = SQLiteShelf("test.sdb", "othertable")
|
||
|
|
||
|
both are empty to start with
|
||
|
|
||
|
>>> d
|
||
|
{}
|
||
|
>>> e
|
||
|
{}
|
||
|
|
||
|
adding stuff is as simple as a regular dict
|
||
|
>>> d['a'] = "moo"
|
||
|
>>> e['a'] = "moo"
|
||
|
|
||
|
regular dict actions work
|
||
|
|
||
|
>>> d['a']
|
||
|
'moo'
|
||
|
>>> e['a']
|
||
|
'moo'
|
||
|
>>> 'a' in d
|
||
|
True
|
||
|
>>> len(d)
|
||
|
1
|
||
|
>>> del d['a']
|
||
|
>>> 'a' in d
|
||
|
False
|
||
|
>>> len(d)
|
||
|
0
|
||
|
>>> del e['a']
|
||
|
|
||
|
objects can be stored in shelves
|
||
|
|
||
|
>> class Test:
|
||
|
.. def __init__(self):
|
||
|
.. self.foo = "bar"
|
||
|
..
|
||
|
>> t = Test()
|
||
|
>> d['t'] = t
|
||
|
>> print d['t'].foo
|
||
|
bar
|
||
|
|
||
|
errors are as normal for a dict
|
||
|
|
||
|
>>> d['x']
|
||
|
Traceback (most recent call last):
|
||
|
...
|
||
|
KeyError: 'x'
|
||
|
>>> del d['x']
|
||
|
Traceback (most recent call last):
|
||
|
...
|
||
|
KeyError: 'x'
|
||
|
|
||
|
Adding and fetching binary strings
|
||
|
|
||
|
>>> d[1] = "a\\x00b"
|
||
|
>>> d[1]
|
||
|
'a\\x00b'
|
||
|
"""
|
||
|
|
||
|
try:
|
||
|
from UserDict import DictMixin
|
||
|
except ImportError:
|
||
|
from collections import MutableMapping as DictMixin
|
||
|
|
||
|
try:
|
||
|
import cPickle as pickle
|
||
|
except ImportError:
|
||
|
import pickle
|
||
|
|
||
|
import sys
|
||
|
valtype = str
|
||
|
if sys.version > '3':
|
||
|
buffer = memoryview
|
||
|
valtype = bytes
|
||
|
|
||
|
import sqlite3
|
||
|
|
||
|
class SQLiteDict(DictMixin):
|
||
|
def __init__(self, filename=':memory:', table='shelf', flags='r', mode=None, valtype=valtype):
|
||
|
self.table = table
|
||
|
self.valtype = valtype
|
||
|
MAKE_SHELF = 'CREATE TABLE IF NOT EXISTS '+self.table+' (key TEXT, value BLOB)'
|
||
|
MAKE_INDEX = 'CREATE UNIQUE INDEX IF NOT EXISTS '+self.table+'_keyndx ON '+self.table+'(key)'
|
||
|
self.conn = sqlite3.connect(filename)
|
||
|
self.conn.text_factory = str
|
||
|
self.conn.execute(MAKE_SHELF)
|
||
|
self.conn.execute(MAKE_INDEX)
|
||
|
self.conn.commit()
|
||
|
|
||
|
def __getitem__(self, key):
|
||
|
GET_ITEM = 'SELECT value FROM '+self.table+' WHERE key = ?'
|
||
|
item = self.conn.execute(GET_ITEM, (key,)).fetchone()
|
||
|
if item is None:
|
||
|
raise KeyError(key)
|
||
|
return self.valtype(item[0])
|
||
|
|
||
|
def __setitem__(self, key, item):
|
||
|
ADD_ITEM = 'REPLACE INTO '+self.table+' (key, value) VALUES (?,?)'
|
||
|
self.conn.execute(ADD_ITEM, (key, sqlite3.Binary(item)))
|
||
|
self.conn.commit()
|
||
|
|
||
|
def __delitem__(self, key):
|
||
|
if key not in self:
|
||
|
raise KeyError(key)
|
||
|
DEL_ITEM = 'DELETE FROM '+self.table+' WHERE key = ?'
|
||
|
self.conn.execute(DEL_ITEM, (key,))
|
||
|
self.conn.commit()
|
||
|
|
||
|
def __iter__(self):
|
||
|
c = self.conn.cursor()
|
||
|
try:
|
||
|
c.execute('SELECT key FROM '+self.table+' ORDER BY key')
|
||
|
for row in c:
|
||
|
yield row[0]
|
||
|
finally:
|
||
|
c.close()
|
||
|
|
||
|
def keys(self):
|
||
|
c = self.conn.cursor()
|
||
|
try:
|
||
|
c.execute('SELECT key FROM '+self.table+' ORDER BY key')
|
||
|
return [row[0] for row in c]
|
||
|
finally:
|
||
|
c.close()
|
||
|
|
||
|
###################################################################
|
||
|
# optional bits
|
||
|
|
||
|
def __len__(self):
|
||
|
GET_LEN = 'SELECT COUNT(*) FROM '+self.table
|
||
|
return self.conn.execute(GET_LEN).fetchone()[0]
|
||
|
|
||
|
def close(self):
|
||
|
if self.conn is not None:
|
||
|
self.conn.commit()
|
||
|
self.conn.close()
|
||
|
self.conn = None
|
||
|
|
||
|
def __del__(self):
|
||
|
self.close()
|
||
|
|
||
|
def __repr__(self):
|
||
|
return repr(dict(self))
|
||
|
|
||
|
class SQLiteShelf(SQLiteDict):
|
||
|
def __getitem__(self, key):
|
||
|
return pickle.loads(SQLiteDict.__getitem__(self, key))
|
||
|
|
||
|
def __setitem__(self, key, item):
|
||
|
SQLiteDict.__setitem__(self, key, pickle.dumps(item))
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
import doctest
|
||
|
doctest.testmod()
|