""" Fast and ugly CLI interface """ import argparse import os import re import sys from sqlalchemy import or_ from pycatalog import scan from pycatalog import misc from pycatalog import dbobjects as dbo from pycatalog.dbcommon import connect, Session from pycatalog import logger BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(30, 38) RESET_SEQ = '\033[0m' COLOR_SEQ = '\033[1;%dm' BOLD_SEQ = '\033[1m' LOG = logger.get_logger(__name__) def colorize(txt, color): """Pretty print with colors to console.""" color_map = {'black': BLACK, 'red': RED, 'green': GREEN, 'yellow': YELLOW, 'blue': BLUE, 'magenta': MAGENTA, 'cyan': CYAN, 'white': WHITE} return COLOR_SEQ % color_map[color] + txt + RESET_SEQ def asserdb(func): def wrapper(args): if not os.path.exists(args.db): print(colorize("File `%s' does not exists!" % args.db, 'red')) sys.exit(1) func(args) return wrapper TYPE_MAP = {0: 'd', 1: 'd', 2: 'f', 3: 'l'} class Iface(object): """Main class which interacts with the pyGTKtalog modules""" def __init__(self, dbname, pretend=False, debug=False): """Init""" self.engine = connect(dbname) self.sess = Session() self.dry_run = pretend self.root = None self._dbname = dbname if debug: scan.LOG.setLevel('DEBUG') LOG.setLevel('DEBUG') def _resolve_path(self, path): """Identify path in the DB""" if not path.startswith("/"): raise AttributeError('Path have to start with slash (/)') last_node = self.root for part in path.split('/'): if not part.strip(): continue for node in last_node.children: if node.filename == part: last_node = node break else: raise AttributeError('No such path: %s' % path) return last_node def _get_full_path(self, file_object): """given the file object, return string with full path to it""" parent = file_object.parent path = [file_object.filename] while parent.type: path.insert(0, parent.filename) parent = parent.parent return u'/' + u'/'.join(path) def _make_path(self, node): """Make the path to the item in the DB""" orig_node = node if node.parent == node: return {u'/': (u' ', 0, u' ')} ext = '' if node.parent.type == dbo.TYPE['root']: ext = colorize(' (%s)' % node.filepath, 'white') path = [] path.append(node.filename) while node.parent != self.root: path.append(node.parent.filename) node = node.parent path = '/'.join([''] + path[::-1]) + ext return {path: (TYPE_MAP[orig_node.type], orig_node.size, orig_node.date)} def _walk(self, dirnode): """Recursively go through the leaves of the node""" items = {} for node in dirnode.children: if node.type == dbo.TYPE['dir']: items.update(self._walk(node)) items.update(self._make_path(node)) return items def _list(self, node): """List only current node content""" items = {} for node in node.children: if node != self.root: items.update(self._make_path(node)) return items def close(self): """Close the session""" self.sess.commit() self.sess.close() def list(self, path=None, recursive=False, long_=False): """Simulate ls command for the provided item path""" self.root = self.sess.query(dbo.File) self.root = self.root.filter(dbo.File.type == dbo.TYPE['root']).first() if path: node = self._resolve_path(path) msg = "Content of path `%s':" % path else: node = self.root msg = "Content of path `/':" print(colorize(msg, 'white')) if recursive: items = self._walk(node) else: items = self._list(node) if long_: filenames = [] format_str = (u'{} {:>%d,} {} {}' % _get_highest_size_length(items)) for fname in sorted(items.keys()): type_, size, date = items[fname] filenames.append(format_str.format(type_, size, date, fname)) else: filenames = sorted(items.keys()) print('\n'.join(filenames)) def update(self, path, dir_to_update=None): """ Update the DB against provided path and optionally directory on the real filesystem """ self.root = self.sess.query(dbo.File) self.root = self.root.filter(dbo.File.type == dbo.TYPE['root']).first() node = self._resolve_path(path) if node == self.root: print(colorize('Cannot update entire db, since root was provided ' 'as path.', 'red')) return if not dir_to_update: dir_to_update = os.path.join(node.filepath, node.filename) if not os.path.exists(dir_to_update): raise OSError("Path to updtate doesn't exists: %s", dir_to_update) print(colorize("Updating node `%s' against directory " "`%s'" % (path, dir_to_update), 'white')) if not self.dry_run: scanob = scan.Scan(dir_to_update) # scanob.update_files(node.id) scanob.update_files(node.id, self.engine) def create(self, dir_to_add, data_dir): """Create new database""" self.root = dbo.File() self.root.id = 1 self.root.filename = 'root' self.root.size = 0 self.root.source = 0 self.root.type = 0 self.root.parent_id = 1 config = dbo.Config() config.key = 'image_path' config.value = data_dir if not self.dry_run: self.sess.add(self.root) self.sess.add(config) self.sess.commit() print(colorize("Creating new db against directory `%s'" % dir_to_add, 'white')) if not self.dry_run: if data_dir == ':same_as_db:': misc.calculate_image_path(None, True) else: misc.calculate_image_path(data_dir, True) scanob = scan.Scan(dir_to_add) scanob.add_files(self.engine) def add(self, dir_to_add): """Add new directory to the db""" self.root = self.sess.query(dbo.File) self.root = self.root.filter(dbo.File.type == 0).first() if not os.path.exists(dir_to_add): raise OSError("Path to add doesn't exists: %s", dir_to_add) print(colorize("Adding directory `%s'" % dir_to_add, 'white')) if not self.dry_run: scanob = scan.Scan(dir_to_add) scanob.add_files() def _annotate(self, item, search_words): """ Find ranges to be highlighted in item, provide them and return result string """ indexes = [] for word in search_words: for match in re.finditer(re.escape(word.lower()), item.lower()): for index in range(match.start(), match.end()): indexes.append(index) highlight = False result = [] for idx, char in enumerate(item): if idx in indexes: if not highlight: highlight = True result.append(COLOR_SEQ % WHITE) result.append(char) else: if highlight: highlight = False result.append(RESET_SEQ) result.append(char) return "".join(result) def find(self, search_words): query = self.sess.query(dbo.File).filter(or_(dbo.File.type == 2, dbo.File.type == 3)) result = [] for word in search_words: phrase = u'%%%s%%' % word query = query.filter(dbo.File.filename.like(phrase)) for item in query.all(): result.append(self._get_full_path(item)) if not result: print("No results for `%s'" % ' '.join(search_words)) return result.sort() for item in result: print(self._annotate(item, search_words)) def fsck(self): """Fsck orphaned images/thumbs""" image_path = (self.sess.query(dbo.Config) .filter(dbo.Config.key == 'image_path')).one().value if image_path == ':same_as_db:': image_path = misc.calculate_image_path(None, False) files_to_remove = [] # remove images/thumbnails which doesn't have file relation for name, obj in (("images", dbo.Image), ("thumbnails", dbo.Thumbnail)): self._purge_orphaned_objects(obj, "Scanning %s " % name) # find all image files not associate with either Image (image/thumb) # or Thumbnail (thumb) objects sys.stdout.write(40 * " " + "\r") count = 0 for root, dirs, files in os.walk(image_path): for fname in files: sys.stdout.write("Scanning files " + "| / - \\".split()[count % 4] + "\r") sys.stdout.flush() count += 1 fname_ = os.path.join(root.split(image_path)[1], fname).lstrip('/') if '_t' in fname: obj = (self.sess.query(dbo.Thumbnail) .filter(dbo.Thumbnail.filename == fname_)).all() if obj: continue obj = (self.sess.query(dbo.Image) .filter(dbo.Image.filename == fname_.replace('_t.', '.'))).all() if obj: continue else: obj = (self.sess.query(dbo.Image) .filter(dbo.Image.filename == fname_)).all() if obj: continue files_to_remove.append(os.path.join(root, fname)) LOG.debug("Found %d orphaned files", len(files_to_remove)) sys.stdout.write(40 * " " + "\r") sys.stdout.flush() if self.dry_run: print("Following files are not associated to any items in the DB:") for filename in sorted(files_to_remove): print(filename) self.sess.rollback() else: _remove_files(image_path, files_to_remove) self.sess.commit() def _purge_orphaned_objects(self, sa_class, msg): """Return tuple of lists of images that are orphaned""" ids_to_remove = [] for count, item in enumerate(self.sess.query(sa_class).all()): sys.stdout.write(msg + "| / - \\".split()[count % 4] + "\r") if not item.file: self.sess.delete(item) ids_to_remove.append(item.id) del item sys.stdout.flush() LOG.debug("Found %d orphaned object of class %s", len(ids_to_remove), sa_class.__name__) self.sess.flush() def _remove_files(image_path, filenames): """Remove files and empty directories in provided location""" count = 0 for count, fname in enumerate(filenames, start=1): os.unlink(fname) LOG.info("Removed %d orphaned files", count) count = 0 for root, dirs, _ in os.walk(image_path): for dirname in dirs: try: os.rmdir(os.path.join(root, dirname)) count += 1 except OSError: pass LOG.info("Removed %d empty directories", count) def _get_highest_size_length(item_dict): highest = len(str(sorted([i[1] for i in item_dict.values()])[-1])) return highest + highest / 3 @asserdb def list_db(args): """List""" obj = Iface(args.db, False, args.debug) obj.list(path=args.path, recursive=args.recursive, long_=args.long) obj.close() @asserdb def update_db(args): """Update""" obj = Iface(args.db, args.pretend, args.debug) obj.update(args.path, dir_to_update=args.dir_to_update) obj.close() @asserdb def add_dir(args): """Add""" obj = Iface(args.db, args.pretend, args.debug) obj.add(args.dir_to_add) obj.close() def create_db(args): """List""" __import__('pdb').set_trace() obj = Iface(args.db, args.pretend, args.debug) obj.create(args.dir_to_add, args.imagedir) obj.close() @asserdb def search(args): """Find""" obj = Iface(args.db, False, args.debug) obj.find(args.search_words) obj.close() @asserdb def cleanup(args): """Cleanup""" obj = Iface(args.db, False, args.debug) obj.fsck() obj.close() def main(): """Main""" parser = argparse.ArgumentParser() subparser = parser.add_subparsers() list_ = subparser.add_parser('list') list_.add_argument('db') list_.add_argument('path', nargs='?') list_.add_argument('-l', '--long', help='Show size, date and type', action='store_true', default=False) list_.add_argument('-r', '--recursive', help='list items in ' 'subdirectories', action='store_true', default=False) list_.add_argument('-d', '--debug', help='Turn on debug', action='store_true', default=False) list_.set_defaults(func=list_db) update = subparser.add_parser('update') update.add_argument('db') update.add_argument('path') update.add_argument('dir_to_update', nargs='?') update.add_argument('-p', '--pretend', help="Don't do the action, just " "give the info what would gonna to happen.", action='store_true', default=False) update.add_argument('-d', '--debug', help='Turn on debug', action='store_true', default=False) update.set_defaults(func=update_db) create = subparser.add_parser('create') create.add_argument('db') create.add_argument('dir_to_add') create.add_argument('-i', '--imagedir', help="Directory where to put " "images for the database. Popular, but deprecated " "choice is `~/.pygtktalog/images'. Current default " "is special string `:same_as_db:' which will try to " "create directory with the same name as the db with " "data suffix", default=':same_as_db:') create.add_argument('-p', '--pretend', help="Don't do the action, just " "give the info what would gonna to happen.", action='store_true', default=False) create.add_argument('-d', '--debug', help='Turn on debug', action='store_true', default=False) create.set_defaults(func=create_db) add = subparser.add_parser('add') add.add_argument('db') add.add_argument('dir_to_add') add.add_argument('-p', '--pretend', help="Don't do the action, just " "give the info what would gonna to happen.", action='store_true', default=False) add.add_argument('-d', '--debug', help='Turn on debug', action='store_true', default=False) add.set_defaults(func=add_dir) find = subparser.add_parser('find') find.add_argument('db') find.add_argument('search_words', nargs='+') find.add_argument('-d', '--debug', help='Turn on debug', action='store_true', default=False) find.set_defaults(func=search) fsck = subparser.add_parser('fsck') fsck.add_argument('db') fsck.add_argument('-p', '--pretend', help="Don't do the action, just give" " the info what would gonna to happen.", action='store_true', default=False) fsck.add_argument('-d', '--debug', help='Turn on debug', action='store_true', default=False) fsck.set_defaults(func=cleanup) args = parser.parse_args() if 'func' in args: args.func(args) else: parser.print_help() if __name__ == '__main__': main()