openmedialibrary/oml/item/scan.py

# -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4


from datetime import datetime
import os
import shutil
import stat
import time

import ox

from changelog import add_record
from item.models import File, Item
from user.models import List
from utils import remove_empty_folders
from websocket import trigger_event
import db
import media
import settings
import state

import logging
logger = logging.getLogger(__name__)

extensions = ['epub', 'pdf', 'txt', 'cbr', 'cbz']

def remove_missing(books=None):
    dirty = False
    logger.debug('remove missing')
    prefix = get_prefix()
    if books is None:
        books = collect_books(prefix)
    with db.session():
        if os.path.exists(prefix):
            logger.debug('scan for removed files')
            db_paths = []
            items = {}
            for f in File.query:
                if state.shutdown:
                    return
                path = f.fullpath()
                db_paths.append(path)
                if f.item:
                    items[path] = f.sha1
                else:
                    state.db.session.delete(f)
                    dirty = True
            if dirty:
                state.db.session.commit()
                dirty = False
            removed = set(db_paths) - set(books)
            if removed:
                logger.debug('%s files removed', len(removed))
                ids = [items[path] for path in removed]
                orphaned = set(ids)
                for i in Item.query.filter(Item.id.in_(ids)):
                    i.remove_file()
                    orphaned.remove(i.id)
                    dirty = True
                if orphaned:
                    logger.debug('%s files orphaned', len(orphaned))
                    for f in File.query.filter(File.sha1.in_(orphaned)):
                        state.db.session.delete(f)
                    dirty = True
            if dirty:
                state.db.session.commit()
                state.cache.clear('group:')
            logger.debug('update filenames')
            for f in File.query:
                if state.shutdown:
                    return
                f.move()
            logger.debug('remove empty folders')
            remove_empty_folders(prefix, True)
    logger.debug('remove missing done')

def add_file(id, f, prefix, from_=None, commit=True):
    user = state.user()
    path = f[len(prefix):]
    logger.debug('%s extract metadata %s', id, path)
    data = media.metadata(f, from_)
    logger.debug('%s create file %s', id, path)
    file = File.get_or_create(id, data, path)
    item = file.item
    item.add_user(user)
    item.added = datetime.utcnow()
    logger.debug('%s load metadata %s', id, path)
    item.load_metadata()
    add_record('additem', item.id, file.info)
    add_record('edititem', item.id, item.meta)
    logger.debug('%s extract icons %s', id, path)
    item.update_icons()
    item.modified = datetime.utcnow()
    logger.debug('%s save item', id)
    item.update(commit=commit)
    logger.debug('%s added', id)
    return file

def get_prefix():
    prefs = settings.preferences
    prefix = os.path.join(os.path.expanduser(prefs['libraryPath']), 'Books' + os.sep)
    if not prefix[-1] == os.sep:
        prefix += os.sep
    assert isinstance(prefix, str)
    return prefix

def collect_books(prefix, status=None):
    logger.debug('collect books')
    books = []
    count = 0
    for root, folders, files in os.walk(prefix):
        for f in files:
            if state.shutdown:
                return []
            if f.startswith('.'):
                continue
            f = os.path.join(root, f)
            ext = f.split('.')[-1].lower()
            if ext == 'kepub':
                ext = 'epub'
            if ext in extensions:
                books.append(os.path.normpath(f))
                count += 1
            if status and not status(count):
                return None
    logger.debug('found %s books', len(books))
    return books

def run_scan():
    logger.debug('run_scan')
    prefix = get_prefix()
    books = collect_books(prefix)
    remove_missing(books)

    added = 0
    with db.session():
        for f in ox.sorted_strings(books):
            if state.shutdown:
                break
            if os.path.exists(f):
                id = media.get_id(f)
                file = File.get(id)
                if not file:
                    file = add_file(id, f, prefix, f)
                    added += 1
    if added:
        trigger_event('change', {})
        logger.debug('imported %s unknown books', added)

def change_path(old, new):
    new_books = os.path.join(new, 'Books')
    if not os.path.exists(new_books):
        ox.makedirs(new)
        shutil.move(os.path.join(old, 'Books'), new_books)
        remove_empty_folders(old)
    else:
        ox.makedirs(new_books)
    run_scan()
    trigger_event('change', {})

def run_import(options=None):
    options = options or {}
    logger.debug('run_import')
    if state.activity.get('cancel'):
        logger.debug('import canceled')
        state.activity = {}
        return
    state.activity = {}
    prefs = settings.preferences
    prefix = os.path.expanduser(options.get('path', prefs['importPath']))
    if os.path.islink(prefix):
        prefix = os.path.realpath(prefix)
    if not prefix[-1] == os.sep:
        prefix += os.sep
    prefix_books = get_prefix()
    prefix_imported = os.path.join(prefix_books, '.import' + os.sep)
    if prefix_books.startswith(prefix) or prefix.startswith(prefix_books):
        error = 'invalid path'
    elif not os.path.exists(prefix):
        error = 'path not found'
    elif not os.path.isdir(prefix):
        error = 'path must be a folder'
    else:
        error = None
    if error:
        trigger_event('activity', {
            'activity': 'import',
            'progress': [0, 0],
            'status': {'code': 404, 'text': error}
        })
        state.activity = {}
        return
    listname = options.get('list')
    if listname:
        listitems = []
    assert isinstance(prefix, str)
    books = []

    def activity(count):
        if count % 100 == 0:
            state.activity = {
                'activity': 'import',
                'path': prefix,
                'progress': [0, count],
            }
            trigger_event('activity', state.activity)
        if state.activity.get('cancel'):
            logger.debug('active import canceled')
            state.activity = {}
            return False
        return True
    books = collect_books(prefix, status=activity)
    if books is None:
        return
    state.activity = {
        'activity': 'import',
        'path': prefix,
        'progress': [0, len(books)],
    }
    trigger_event('activity', state.activity)
    position = 0
    added = 0
    last = 0
    for f in ox.sorted_strings(books):
        position += 1
        if not os.path.exists(f):
            continue
        with db.session():
            id = media.get_id(f)
            file = File.get(id)
            f_import = f
            if not file:
                f = f.replace(prefix, prefix_imported)
                ox.makedirs(os.path.dirname(f))
                if options.get('mode') == 'move':
                    try:
                        shutil.move(f_import, f)
                    except:
                        shutil.copy2(f_import, f)
                else:
                    shutil.copy2(f_import, f)
                file = add_file(id, f, prefix_books, f_import)
                file.move()
                added += 1
            elif options.get('mode') == 'move':
                try:
                    os.chmod(f_import, stat.S_IWRITE)
                    os.unlink(f_import)
                except:
                    pass
            if listname:
                listitems.append(file.item.id)
        if state.activity.get('cancel'):
            state.activity = {}
            return
        if state.shutdown:
            return
        if time.time() - last > 5:
            last = time.time()
            state.activity = {
                'activity': 'import',
                'progress': [position, len(books)],
                'path': prefix,
                'added': added,
            }
            trigger_event('activity', state.activity)

    if listname and listitems:
        with db.session():
            l = List.get(settings.USER_ID, listname)
            if l:
                l.add_items(listitems)
    trigger_event('activity', {
        'activity': 'import',
        'progress': [position, len(books)],
        'path': prefix,
        'status': {'code': 200, 'text': ''},
        'added': added,
    })
    state.activity = {}
    remove_empty_folders(prefix_books)
    if options.get('mode') == 'move':
        remove_empty_folders(prefix, True)

def import_folder():
    if not (state.activity and state.activity.get('activity') == 'import'):
        import_path = settings.preferences['importPath']
        import_path = os.path.normpath(os.path.expanduser(import_path))
        import_path_base = os.path.dirname(import_path)
        if not os.path.exists(import_path) and os.path.exists(import_path_base):
            os.makedirs(import_path)
        logger.debug('scan importPath %s', import_path)
        if os.path.exists(import_path):
            run_import({
                'path': import_path,
                'mode': 'move'
            })
            remove_empty_folders(import_path, True)
    if state.main:
        state.main.call_later(10*60, lambda: state.tasks.queue('scanimport'))