CouchPotatoServer/couchpotato/core/plugins/library/main.py

from couchpotato import get_session
from couchpotato.core.event import addEvent, fireEventAsync, fireEvent
from couchpotato.core.helpers.encoding import toUnicode, simplifyString
from couchpotato.core.logger import CPLog
from couchpotato.core.plugins.base import Plugin
from couchpotato.core.settings.model import Library, LibraryTitle, File
from string import ascii_letters
import traceback

log = CPLog(__name__)

class LibraryPlugin(Plugin):

    default_dict = {'titles': {}, 'files':{}}

    def __init__(self):
        addEvent('library.add', self.add)
        addEvent('library.update', self.update)

    def add(self, attrs = {}, update_after = True):

        db = get_session()

        l = db.query(Library).filter_by(identifier = attrs.get('identifier')).first()
        if not l:
            status = fireEvent('status.get', 'needs_update', single = True)
            l = Library(
                year = attrs.get('year'),
                identifier = attrs.get('identifier'),
                plot = toUnicode(attrs.get('plot')),
                tagline = toUnicode(attrs.get('tagline')),
                status_id = status.get('id')
            )

            title = LibraryTitle(
                title = toUnicode(attrs.get('title')),
                simple_title = self.simplifyTitle(attrs.get('title'))
            )

            l.titles.append(title)

            db.add(l)
            db.commit()

        # Update library info
        if update_after is not False:
            handle = fireEventAsync if update_after is 'async' else fireEvent
            handle('library.update', identifier = l.identifier, default_title = attrs.get('title', ''))

        return l.to_dict(self.default_dict)

    def update(self, identifier, default_title = '', force = False):

        db = get_session()
        library = db.query(Library).filter_by(identifier = identifier).first()
        done_status = fireEvent('status.get', 'done', single = True)

        if library:
            library_dict = library.to_dict(self.default_dict)

        do_update = True

        if library.status_id == done_status.get('id') and not force:
            do_update = False
        else:
            info = fireEvent('movie.info', merge = True, identifier = identifier)
            if not info or len(info) == 0:
                log.error('Could not update, no movie info to work with: %s' % identifier)
                return False

        # Main info
        if do_update:
            library.plot = toUnicode(info.get('plot', ''))
            library.tagline = toUnicode(info.get('tagline', ''))
            library.year = info.get('year', 0)
            library.status_id = done_status.get('id')
            library.info = info
            db.commit()

            # Titles
            [db.delete(title) for title in library.titles]
            db.commit()

            titles = info.get('titles', [])
            log.debug('Adding titles: %s' % titles)
            for title in titles:
                t = LibraryTitle(
                    title = toUnicode(title),
                    simple_title = self.simplifyTitle(title),
                    default = title.lower() == default_title.lower() or (default_title is '' and titles[0] == title)
                )
                library.titles.append(t)

            db.commit()

            # Files
            images = info.get('images', [])
            for type in images:
                for image in images[type]:
                    if not isinstance(image, str):
                        continue

                    file_path = fireEvent('file.download', url = image, single = True)
                    file_obj = fireEvent('file.add', path = file_path, type_tuple = ('image', type), single = True)
                    try:
                        file_obj = db.query(File).filter_by(id = file_obj.get('id')).one()
                        library.files.append(file_obj)
                        db.commit()
                    except:
                        log.debug('Failed to attach to library: %s' % traceback.format_exc())

            library_dict = library.to_dict(self.default_dict)

        fireEvent('library.update_finish', data = library_dict)

        return library_dict

    def simplifyTitle(self, title):

        title = toUnicode(title)
        nr_prefix = '' if title[0] in ascii_letters else '#'
        title = simplifyString(title)

        for prefix in ['the ']:
            if prefix == title[:len(prefix)]:
                title = title[len(prefix):]
                break

        return nr_prefix + title