import struct
import base64
import shelve
import os.path
import utils, db, helpers, folder, image, log
"""
Methods for generating/using ShortURLs
"""
def int2key (id) :
"""
Turn an integer into a short-as-possible url-safe string
"""
for type in ('B', 'H', 'I') :
try :
return base64.b64encode(struct.pack(type, id), '-_').rstrip('=')
except struct.error :
continue
raise Exception("ID overflow: %s" % id)
def key2int (key) :
# base64 ignores extra padding, but if it doesn't, it's (4 - len%4), if len%4 != 0
# and it breaks on unicode strings
bytes = base64.b64decode(str(key + '='*6), '-_')
type = {
1: 'B',
2: 'H',
4: 'I',
}[len(bytes)]
return struct.unpack(type, bytes)[0]
class DB (object) :
def __init__ (self, read_only=True) :
self.db = shelve.open('shorturls2', read_only and 'r' or 'c')
def html_path (self, key, index) :
type, dirpath, fname = self.db[key]
if type == 'img' :
fname += '.html'
elif type == 'dir' :
fname = ''
if index :
dirpath = '../%s' % dirpath
if type == 'dir' and index > 1 :
fname = 'index_%s.html' % (index - 1)
return os.path.join(dirpath, fname)
def image_info (self, key) :
type, dirpath, fname = self.db[key]
if type != 'img' :
raise ValueError("%s is not an img" % key)
return dirpath, fname
def shorturls_for (self, paths) :
ret = []
for key in self.db.keys() :
if key.startswith('_') :
continue
type, dir, fname = self.db[key]
path = os.path.join(dir.lstrip('.').lstrip('/'), fname)
if path in paths :
ret.append(key)
paths.remove(path)
if paths :
raise ValueError("Paths not found: %s" % " ".join(paths))
return ret
def html_path (key, index=None) :
dir, fname = node_info(key)
if fname :
return utils.url(dir, fname + '.html')
else :
return utils.url(dir, helpers.url_for_page(index or 0))
def node_info (key) :
res = db.select("""SELECT dirpath, filename FROM nodes WHERE id=?""", key2int(key)).fetchone()
if res :
return res
else :
raise KeyError(key)
def image_info (key) :
res = db.select("""SELECT dirpath, filename FROM images WHERE id=?""", key2int(key)).fetchone()
if res :
return res
else :
raise KeyError(key)
def get_images (keys) :
res = [db.select("""SELECT dirpath, filename FROM images WHERE id=?""", key2int(key)).fetchone() for key in keys]
# don't mind if we don't get as many as we asked for?
if res :
return res
else :
raise KeyError(keys)
def _got_obj_key (obj, id) :
key = int2key(id)
obj.shorturl_code = key
if isinstance(obj, folder.Folder) :
dir, fname = utils.strip_path(obj.path), ''
elif isinstance(obj, image.Image) :
dir, fname = utils.strip_path(obj.dir.path), obj.name
else :
assert(False, "%r %r" % (obj, id))
log.info("%6s -> %s/%s", key, dir, fname)
def updateDB (root) :
"""
Update the SQL database
type - one of 'img', 'dir'
dirpath - the path to the directory, e.g. '.', './foobar', './foobar/quux'
fname - the filename, one of '', 'DSC9839.JPG', 'this.png', etc.
"""
dirqueue = [root]
# dict of (dir, fname) -> obj
paths = {}
while dirqueue :
dir = dirqueue.pop(0)
dirqueue.extend(dir.subdirs.itervalues())
if dir.alive :
pathtuple = (utils.strip_path(dir.path), '')
log.debug("dir %50s", pathtuple[0])
paths[pathtuple] = dir
for img in dir.images.itervalues() :
pathtuple = (utils.strip_path(img.dir.path), img.name)
log.debug("img %50s %15s", *pathtuple)
paths[pathtuple] = img
log.info("we have %d nodes", len(paths))
for (id, dir, fname) in db.select("SELECT id, dirpath, filename FROM nodes") :
try :
obj = paths.pop((dir, fname))
key = int2key(id)
obj.shorturl_code = key
log.debug("%s %50s %15s -> %d %s", dir and "img" or "dir", dir, fname, id, key)
except KeyError :
pass
# log.warning("non-existant node (%d, %s, %s) in db", id, dir, fname)
if paths :
log.info("allocating shorturls for %d new nodes:", len(paths))
db.insert_many(
_got_obj_key,
"INSERT INTO nodes (dirpath, filename) VALUES (?, ?)",
((obj, (path, fname)) for ((path, fname), obj) in paths.iteritems())
)
else :
log.info("no new images")