alexandria.py
author Dmitriy Morozov <morozov@cs.duke.edu>
Sat, 24 May 2008 06:23:33 -0400
changeset 6 3f2c0ca0812a
parent 5 35552b6cdc51
child 7 d0e3cd42a3af
permissions -rwxr-xr-x
Command dispatcher notifies of ambiguous choice from the given prefix

#!/usr/bin/env python

import os, sys, os.path
import hashlib
import terminal
from optparse import OptionParser
from models import Author, Paper, Tag, AuthorNickname, initDatabase, session, asc

db_filename = 'alexandria.db'
term = terminal.TerminalController()
color = {'title':   term.GREEN + term.BOLD, 
         'author':  term.YELLOW + term.BOLD,
         'label':   term.CYAN + term.BOLD,
         'path':    term.NORMAL,
         'hash':    term.RED + term.BOLD,
         'error':   term.RED + term.BOLD,
         'normal':  term.NORMAL}
default_viewer = '/usr/bin/evince'

def find_database(starting_path = None):
    if not starting_path: starting_path = '.'

    # Walk up until "alexandria.db" is found
    # return (True, path) if found, (False, os.path.join(starting_path, 'alexandria.db')) otherwise
    directory = starting_path
    while os.path.abspath(directory) != '/':
        if os.path.exists(os.path.join(directory, db_filename)):
            break
        directory = os.path.abspath(os.path.join(directory, '..'))
    else:
        return (False, os.path.join(starting_path, db_filename))
    return (True, os.path.join(directory, db_filename))

def add(args, options):
    path = args[0]
    if not os.path.exists(path):
        print _colorize_string('error', "Path %s does not exist. Cannot add paper" % path)
        return

    m = hashlib.md5()
    fd = open(path, 'r')
    m.update(fd.read())
    fd.close()

    path = os.path.abspath(path)
    p = Paper.get_by(path = path) or Paper.get_by(md5 = m.hexdigest())
    if p is not None:
        print _colorize_string('error', "Paper already exists, use update")
        print '--------------------------------'
        _show_paper(p)
        return

    p = Paper(path = path, md5 = m.hexdigest())
    _set_options(p, options, required = ['title'])

    session.flush()
    _show_paper(p)

def update(args, options):
    p = Paper.query.filter(Paper.md5.startswith(args[0])).one()
    _set_options(p, options)
    session.flush()
    _show_paper(p)

def view(args, options):
    if len(args) < 1: return
    p = Paper.query.filter(Paper.md5.startswith(args[0])).one()
    if not p: return
    if len(args) > 1:   viewer = args[1]
    else:               viewer = default_viewer
    os.system('%s %s' % (viewer, p.path))

def remove(args, options):
    if len(args) < 1: return
    p = Paper.query.filter(Paper.md5.startswith(args[0])).one()
    if not p: return
    print "Removing"
    _show_paper(p)
    p.delete()
    session.flush()

def list(args, options):
    papers = Paper.query

    # Refactor with _set_options()
    for label_with_commas in (options.labels or []):
        labels = label_with_commas.split(',')
        for label in labels:
            label = unicode(label.strip())
            label = label.replace('*', '%')             # allow for glob style-pattern
            papers = papers.filter(Paper.tags.any(Tag.name.like(label)))

    for author_with_commas in (options.authors or []):
        authors = author_with_commas.split(',')
        for author in authors:
            author = unicode(author.strip())
            an = AuthorNickname.get_by(name = author)
            if an: a = an.author.name
            else:  a = author.replace('*', '%')
            papers = papers.filter(Paper.authors.any(Author.name.like(a)))

    print
    for p in papers.all():
        _show_paper(p)
        print

def alias(args, options):
    if len(args) > 1:
        a =  Author.get_by_or_init(name = unicode(args[1]))
        an = AuthorNickname.get_by_or_init(name = unicode(args[0]))
        an.author = a
        session.flush()

    print "Nicknames:"
    for a in Author.query.all():
        if len(a.nicknames) > 0:
            print '  ' + a.name + ':',
            for an in a.nicknames[:-1]:
                print an.name + ',',
            #print '%s: %s' % (a.nicknames[-1], a.name)
            print a.nicknames[-1]

def labels(args, options):
    if len(args) >= 2:
        t = Tag.get_by(name = unicode(args[0]))
        if t:
            t.name = unicode(args[1])
    session.flush()

    print "Labels:"
    for t in Tag.query.order_by(asc(Tag.name)).all():
        if len(t.papers) == 0:                  # clean the database
            t.delete()
            continue
        print '  (%d) %s' % (len(t.papers), t.name)
    session.flush()

def _set_options(p, options, required = []):
    title = options.title or ('title' in required) and raw_input("Enter title: ")
    if title:
        p.title = unicode(title)

    for label_with_commas in (options.labels or []):
        labels = label_with_commas.split(',')
        for label in labels:
            label = unicode(label.strip())
            if label[0] == '-':         # remove label
                t = Tag.get_by(name = label[1:])
                t.papers.remove(p)
            else:                       # add label
                t = Tag.get_by_or_init(name = label)
                t.papers.append(p)

    for author_with_commas in (options.authors or []):
        authors = author_with_commas.split(',')
        for author in authors:
            author = unicode(author.strip())
            an = AuthorNickname.get_by(name = author)
            if an: a = an.author
            else:  a = Author.get_by_or_init(name = author)
            a.papers.append(p)

def _sort_authors(authors):
    authors.sort()          # FIXME: deal with firstname lastname issues

def _show_paper(paper):
    print _colorize_string('title', paper.title)
    authors = [str(a) for a in paper.authors]
    _sort_authors(authors)
    for author in authors[:-1]:
        print '%s,' % _colorize_string('author', author),
    print '%s' % _colorize_string('author', authors[-1])
    print 'Labels:',
    for tag in paper.tags:
        print '+%s' % _colorize_string('label', tag),
    print color['normal']
    print "Path:   %s" % _colorize_string('path', paper.path)
    print "Hash:   %s" % _colorize_string('hash', paper.md5)

def _colorize_string(clr, str):
    return '%s%s%s' % (color[clr], str, color['normal'])

commands = [
            (add,       'add a paper to the database'),
            (list,      'list papers in the database'),
            (alias,     'add or list author nicknames'),
            (update,    'update paper by hash'),
            (view,      'view paper by hash'),
            (remove,    'remove paper by hash'),
            (labels,    'rename and/or list labels')
           ]


if __name__ == "__main__":
    usage =  '%s COMMAND OPTIONS\n' % sys.argv[0]
    usage += 'Commands:\n'
    for cmd in commands:
        func, description = cmd
        usage += '  %-10s - %s\n' % (func.__name__, description)
    
    # Parse options
    parser = OptionParser(usage = usage)
    parser.add_option('-a', '--author', action='append', dest='authors', help='author')
    parser.add_option('-t', '--title', dest='title', help='title')
    parser.add_option('-l', '--label', action='append', dest='labels', help='label')
    parser.add_option('-s', '--hash', dest='hash', help='hash (only for list)')
    parser.add_option('-D', '--database', dest='database', help='directory with the database')
    (options, args) = parser.parse_args()
    
    # Find database
    found, path = find_database(options.database)
    initDatabase(path, not found)
    
    if len(args) == 0: sys.exit()
    candidates = []
    for cmd in commands:
        func, description = cmd
        if func.__name__.startswith(args[0]):
            candidates += [func]
    if len(candidates) > 1:
        print "Ambiguous choices:",
        for c in candidates: print c.__name__,
        print
    else:
        candidates[0](args[1:], options)