1

import struct

     2

import base64

     3

import shelve

     4

import os.path

     5

     7

    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