1 import struct

     2 import base64

     3 import shelve

     4 import os.path

     5 

     6 import utils, db, helpers, folder, image, log

     7 

     8 """

     9     Methods for generating/using ShortURLs

    10 """

    11 

    12 def int2key (id) :

    13     """

    14         Turn an integer into a short-as-possible url-safe string

    15     """

    16     for type in ('B', 'H', 'I') :

    17         try :

    18             return base64.b64encode(struct.pack(type, id), '-_').rstrip('=')

    19         except struct.error :

    20             continue

    21 

    22     raise Exception("ID overflow: %s" % id)

    23 

    24 def key2int (key) :

    25     # base64 ignores extra padding, but if it doesn't, it's (4 - len%4), if len%4 != 0

    26     # and it breaks on unicode strings

    27     bytes = base64.b64decode(str(key + '='*6), '-_')

    28     

    29     type = {

    30         1: 'B',

    31         2: 'H',

    32         4: 'I',

    33     }[len(bytes)]

    34 

    35     return struct.unpack(type, bytes)[0]

    36 

    37 class DB (object) :

    38     def __init__ (self, read_only=True) :

    39         self.db = shelve.open('shorturls2', read_only and 'r' or 'c')

    40 

    41     def html_path (self, key, index) :

    42         type, dirpath, fname = self.db[key]

    43 

    44         if type == 'img' :

    45             fname += '.html'

    46         elif type == 'dir' :

    47             fname = ''

    48 

    49         if index :

    50             dirpath = '../%s' % dirpath

    51             

    52             if type == 'dir' and index > 1 : 

    53                 fname = 'index_%s.html' % (index - 1)

    54 

    55         return os.path.join(dirpath, fname)

    56    

    57     def image_info (self, key) :

    58         type, dirpath, fname = self.db[key]

    59 

    60         if type != 'img' :

    61             raise ValueError("%s is not an img" % key)

    62 

    63         return dirpath, fname

    64     

    65     def shorturls_for (self, paths) :

    66         ret = []

    67 

    68         for key in self.db.keys() :

    69             if key.startswith('_') :

    70                 continue

    71 

    72             type, dir, fname = self.db[key]

    73             path = os.path.join(dir.lstrip('.').lstrip('/'), fname) 

    74             if path in paths :

    75                 ret.append(key)

    76                 paths.remove(path)

    77         

    78         if paths :

    79             raise ValueError("Paths not found: %s" % " ".join(paths))

    80 

    81         return ret

    82 

    83 def html_path (key, index=None) :

    84     dir, fname = node_info(key)

    85 

    86     if fname :

    87         return utils.url(dir, fname + '.html')

    88     else :

    89         return utils.url(dir, helpers.url_for_page(index or 0))

    90 

    91 def node_info (key) :

    92     res = db.select("""SELECT dirpath, filename FROM nodes WHERE id=?""", key2int(key)).fetchone()

    93     

    94     if res :

    95         return res

    96 

    97     else :

    98         raise KeyError(key)

    99 

   100 def image_info (key) :

   101     res = db.select("""SELECT dirpath, filename FROM images WHERE id=?""", key2int(key)).fetchone()

   102     

   103     if res :

   104         return res

   105 

   106     else :

   107         raise KeyError(key)

   108    

   109 def get_images (keys) :

   110     res = [db.select("""SELECT dirpath, filename FROM images WHERE id=?""", key2int(key)).fetchone() for key in keys]

   111 

   112     # don't mind if we don't get as many as we asked for?

   113     if res :

   114         return res

   115 

   116     else :

   117         raise KeyError(keys)

   118 

   119 def _got_obj_key (obj, id) :

   120     key = int2key(id)

   121 

   122     obj.shorturl_code = key

   123 

   124     if isinstance(obj, folder.Folder) :

   125         dir, fname = utils.strip_path(obj.path), ''

   126     elif isinstance(obj, image.Image) :

   127         dir, fname = utils.strip_path(obj.dir.path), obj.name

   128     else :

   129         assert(False, "%r %r" % (obj, id))

   130 

   131     log.info("%6s -> %s/%s", key, dir, fname)

   132 

   133 def updateDB (root) :

   134     """

   135         Update the SQL database

   136 

   137         type    - one of 'img', 'dir'

   138         dirpath - the path to the directory, e.g. '.', './foobar', './foobar/quux'

   139         fname   - the filename, one of '', 'DSC9839.JPG', 'this.png', etc.

   140     """

   141 

   142     dirqueue = [root]

   143 

   144     # dict of (dir, fname) -> obj

   145     paths = {}

   146 

   147     while dirqueue :

   148         dir = dirqueue.pop(0)

   149 

   150         dirqueue.extend(dir.subdirs.itervalues())

   151 

   152         if dir.alive :

   153             pathtuple = (utils.strip_path(dir.path), '')

   154             

   155             log.debug("dir %50s", pathtuple[0])

   156 

   157             paths[pathtuple] = dir

   158 

   159         for img in dir.images.itervalues() :

   160             pathtuple = (utils.strip_path(img.dir.path), img.name)

   161             

   162             log.debug("img %50s %15s", *pathtuple)

   163 

   164             paths[pathtuple] = img

   165     

   166     log.info("we have %d nodes", len(paths))

   167 

   168     for (id, dir, fname) in db.select("SELECT id, dirpath, filename FROM nodes") :

   169         try :

   170             obj = paths.pop((dir, fname))

   171             key = int2key(id)

   172 

   173             obj.shorturl_code = key

   174 

   175             log.debug("%s %50s %15s -> %d %s", dir and "img" or "dir", dir, fname, id, key)

   176         

   177         except KeyError :

   178             pass

   179 #            log.warning("non-existant node (%d, %s, %s) in db", id, dir, fname)

   180     

   181     if paths :

   182         log.info("allocating shorturls for %d new nodes:", len(paths))

   183 

   184         db.insert_many(

   185             _got_obj_key,

   186             "INSERT INTO nodes (dirpath, filename) VALUES (?, ?)",

   187             ((obj, (path, fname)) for ((path, fname), obj) in paths.iteritems())

   188         )

   189     else :

   190         log.info("no new images")

   191