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

from couchpotato import get_session
from couchpotato.api import addApiView
from couchpotato.core.event import addEvent
from couchpotato.core.helpers.encoding import toUnicode
from couchpotato.core.helpers.variable import mergeDicts, md5, getExt
from couchpotato.core.logger import CPLog
from couchpotato.core.plugins.base import Plugin
from couchpotato.core.settings.model import Quality, Profile, ProfileType
from sqlalchemy.sql.expression import or_
import re
import time

log = CPLog(__name__)


class QualityPlugin(Plugin):

    qualities = [
        {'identifier': 'bd50', 'hd': True, 'size': (15000, 60000), 'label': 'BR-Disk', 'alternative': ['bd25'], 'allow': ['1080p'], 'ext':[], 'tags': ['bdmv', 'certificate', ('complete', 'bluray')]},
        {'identifier': '1080p', 'hd': True, 'size': (4000, 20000), 'label': '1080p', 'width': 1920, 'height': 1080, 'alternative': [], 'allow': [], 'ext':['mkv', 'm2ts'], 'tags': ['m2ts']},
        {'identifier': '720p', 'hd': True, 'size': (3000, 10000), 'label': '720p', 'width': 1280, 'height': 720, 'alternative': [], 'allow': [], 'ext':['mkv', 'ts']},
        {'identifier': 'brrip', 'hd': True, 'size': (700, 7000), 'label': 'BR-Rip', 'alternative': ['bdrip'], 'allow': ['720p', '1080p'], 'ext':['avi'], 'tags': ['hdtv', 'hdrip', 'webdl', ('web', 'dl')]},
        {'identifier': 'dvdr', 'size': (3000, 10000), 'label': 'DVD-R', 'alternative': [], 'allow': [], 'ext':['iso', 'img'], 'tags': ['pal', 'ntsc', 'video_ts', 'audio_ts']},
        {'identifier': 'dvdrip', 'size': (600, 2400), 'label': 'DVD-Rip', 'width': 720, 'alternative': [], 'allow': [], 'ext':['avi', 'mpg', 'mpeg'], 'tags': [('dvd', 'rip'), ('dvd', 'xvid'), ('dvd', 'divx')]},
        {'identifier': 'scr', 'size': (600, 1600), 'label': 'Screener', 'alternative': ['screener', 'dvdscr', 'ppvrip', 'dvdscreener'], 'allow': ['dvdr', 'dvd'], 'ext':['avi', 'mpg', 'mpeg'], 'tags': ['webrip', ('web', 'rip')]},
        {'identifier': 'r5', 'size': (600, 1000), 'label': 'R5', 'alternative': ['r6'], 'allow': ['dvdr'], 'ext':['avi', 'mpg', 'mpeg']},
        {'identifier': 'tc', 'size': (600, 1000), 'label': 'TeleCine', 'alternative': ['telecine'], 'allow': [], 'ext':['avi', 'mpg', 'mpeg']},
        {'identifier': 'ts', 'size': (600, 1000), 'label': 'TeleSync', 'alternative': ['telesync', 'hdts'], 'allow': [], 'ext':['avi', 'mpg', 'mpeg']},
        {'identifier': 'cam', 'size': (600, 1000), 'label': 'Cam', 'alternative': ['camrip', 'hdcam'], 'allow': [], 'ext':['avi', 'mpg', 'mpeg']}
    ]
    pre_releases = ['cam', 'ts', 'tc', 'r5', 'scr']

    def __init__(self):
        addEvent('quality.all', self.all)
        addEvent('quality.single', self.single)
        addEvent('quality.guess', self.guess)
        addEvent('quality.pre_releases', self.preReleases)

        addApiView('quality.size.save', self.saveSize)
        addApiView('quality.list', self.allView, docs = {
            'desc': 'List all available qualities',
            'return': {'type': 'object', 'example': """{
            'success': True,
            'list': array, qualities
}"""}
        })

        addEvent('app.initialize', self.fill, priority = 10)

    def preReleases(self):
        return self.pre_releases

    def allView(self, **kwargs):

        return {
            'success': True,
            'list': self.all()
        }

    def all(self):

        db = get_session()

        qualities = db.query(Quality).all()

        temp = []
        for quality in qualities:
            q = mergeDicts(self.getQuality(quality.identifier), quality.to_dict())
            temp.append(q)

        return temp

    def single(self, identifier = ''):

        db = get_session()
        quality_dict = {}

        quality = db.query(Quality).filter(or_(Quality.identifier == identifier, Quality.id == identifier)).first()
        if quality:
            quality_dict = dict(self.getQuality(quality.identifier), **quality.to_dict())

        return quality_dict

    def getQuality(self, identifier):

        for q in self.qualities:
            if identifier == q.get('identifier'):
                return q

    def saveSize(self, **kwargs):

        db = get_session()
        quality = db.query(Quality).filter_by(identifier = kwargs.get('identifier')).first()

        if quality:
            setattr(quality, kwargs.get('value_type'), kwargs.get('value'))
            db.commit()

        return {
            'success': True
        }

    def fill(self):

        db = get_session();

        order = 0
        for q in self.qualities:

            # Create quality
            qual = db.query(Quality).filter_by(identifier = q.get('identifier')).first()

            if not qual:
                log.info('Creating quality: %s', q.get('label'))
                qual = Quality()
                qual.order = order
                qual.identifier = q.get('identifier')
                qual.label = toUnicode(q.get('label'))
                qual.size_min, qual.size_max = q.get('size')

                db.add(qual)

            # Create single quality profile
            prof = db.query(Profile).filter(
                    Profile.core == True
                ).filter(
                    Profile.types.any(quality = qual)
                ).all()

            if not prof:
                log.info('Creating profile: %s', q.get('label'))
                prof = Profile(
                    core = True,
                    label = toUnicode(qual.label),
                    order = order
                )
                db.add(prof)

                profile_type = ProfileType(
                    quality = qual,
                    profile = prof,
                    finish = True,
                    order = 0
                )
                prof.types.append(profile_type)

            order += 1

        db.commit()

        time.sleep(0.3) # Wait a moment

        return True

    def guess(self, files, extra = {}):

        # Create hash for cache
        hash = md5(str([f.replace('.' + getExt(f), '') for f in files]))
        cached = self.getCache(hash)
        if cached and extra is {}: return cached

        for cur_file in files:
            words = re.split('\W+', cur_file.lower())

            found = {}
            for quality in self.all():
                contains = self.containsTag(quality, words, cur_file)
                if contains:
                    found[quality['identifier']] = True

            if 'web' in words and 'dl' in words:
                print found

            for quality in self.all():

                # Check identifier
                if quality['identifier'] in words:
                    if len(found) == 0 or len(found) == 1 and found.get(quality['identifier']):
                        log.debug('Found via identifier "%s" in %s', (quality['identifier'], cur_file))
                        return self.setCache(hash, quality)

                # Check alt and tags
                contains = self.containsTag(quality, words, cur_file)
                if contains:
                    return self.setCache(hash, quality)

        # Try again with loose testing
        quality = self.guessLoose(hash, files = files, extra = extra)
        if quality:
            return self.setCache(hash, quality)

        log.debug('Could not identify quality for: %s', files)
        return None

    def containsTag(self, quality, words, cur_file = ''):

        # Check alt and tags
        for tag_type in ['alternative', 'tags']:
            for alt in quality.get(tag_type, []):
                if isinstance(alt, tuple) and '.'.join(alt) in '.'.join(words):
                    log.debug('Found %s via %s %s in %s', (quality['identifier'], tag_type, quality.get(tag_type), cur_file))
                    return True

            if list(set(quality.get(tag_type, [])) & set(words)):
                log.debug('Found %s via %s %s in %s', (quality['identifier'], tag_type, quality.get(tag_type), cur_file))
                return True

        return

    def guessLoose(self, hash, files = None, extra = None):

        if extra:
            for quality in self.all():

                # Check width resolution, range 20
                if quality.get('width') and (quality.get('width') - 20) <= extra.get('resolution_width', 0) <= (quality.get('width') + 20):
                    log.debug('Found %s via resolution_width: %s == %s', (quality['identifier'], quality.get('width'), extra.get('resolution_width', 0)))
                    return self.setCache(hash, quality)

                # Check height resolution, range 20
                if quality.get('height') and (quality.get('height') - 20) <= extra.get('resolution_height', 0) <= (quality.get('height') + 20):
                    log.debug('Found %s via resolution_height: %s == %s', (quality['identifier'], quality.get('height'), extra.get('resolution_height', 0)))
                    return self.setCache(hash, quality)

            if 480 <= extra.get('resolution_width', 0) <= 720:
                log.debug('Found as dvdrip')
                return self.setCache(hash, self.single('dvdrip'))

        return None