<commit_before>import logging
import datetime
from flexget.manager import Session
from flexget.plugin import register_plugin, PluginError, get_plugin_by_name, priority, register_parser_option
from flexget.manager import Base
from sqlalchemy import Column, Integer, String, DateTime, Boolean, Unicode
from flexget.utils.imdb import extract_id, ImdbSearch, ImdbParser, log as imdb_log
from flexget.utils.tools import str_to_boolean

log = logging.getLogger('imdb_queue')


class ImdbQueue(Base):

    __tablename__ = 'imdb_queue'

    id = Column(Integer, primary_key=True)
    imdb_id = Column(String)
    quality = Column(String)
    title = Column(Unicode)
    immortal = Column(Boolean)
    added = Column(DateTime)

    def __init__(self, imdb_id, quality, immortal):
        self.imdb_id = imdb_id
        self.quality = quality
        self.immortal = immortal
        self.added = datetime.datetime.now()

    def __str__(self):
        return '<ImdbQueue(imdb_id=%s,quality=%s,force=%s)>' % (self.imdb_id, self.quality, self.immortal)


class FilterImdbQueue(object):
    """
    Allows a queue of upcoming movies that will be forcibly allowed if the given quality matches

    Example:

    imdb_queue: yes
    """

    def validator(self):
        from flexget import validator
        return validator.factory('boolean')

    @priority(129)
    def on_feed_filter(self, feed):
        for entry in feed.entries:
            # make sure the entry has IMDB fields filled
            try:
                get_plugin_by_name('imdb_lookup').instance.lookup(feed, entry)
            except PluginError:
                # no IMDB data, can't do anything
                continue

            # entry has imdb url
            if 'imdb_url' in entry:
                # get an imdb id
                if 'imdb_id' in entry and entry['imdb_id'] is not None:
                    imdb_id = entry['imdb_id']
                else:
                    imdb_id = extract_id(entry['imdb_url'])

                if not imdb_id:
                    log.warning("No imdb id could be determined for %s" % entry['title'])
                    continue

                if not 'quality' in entry:
                    log.warning('No quality found for %s, assigning unknown.' % entry['title'])
                    entry['quality'] = 'unknown'

                item = feed.session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).\
                                                     filter((ImdbQueue.quality == entry['quality']) | (ImdbQueue.quality == "ANY")).first()

                if item:
                    entry['immortal'] = item.immortal
                    log.info("Accepting %s from queue with quality %s. Force: %s" % (entry['title'], entry['quality'], entry['immortal']))
                    feed.accept(entry, 'imdb-queue - force: %s' % entry['immortal'])
                    # and remove from database
                    feed.session.delete(item)
                else:
                    log.log(5, "%s not in queue with wanted quality, skipping" % entry['title'])


class ImdbQueueManager(object):
    """
    Handle IMDb queue management; add, delete and list
    """

    valid_actions = ['add', 'del', 'list']

    options = {}

    @staticmethod
    def optik_imdb_queue(option, opt, value, parser):
        """
        Callback for Optik
        --imdb-queue (add|del|list) [IMDB_URL|NAME] [quality]
        """
        if len(parser.rargs) == 0:
            print 'Usage: --imdb-queue (add|del|list) [IMDB_URL|NAME] [QUALITY] [FORCE]'
            # set some usage option so that feeds will be disabled later
            ImdbQueueManager.options['usage'] = True
            return

        ImdbQueueManager.options['action'] = parser.rargs[0].lower()

        if len(parser.rargs) == 1:
            return
        # 2 args is the minimum allowed (operation + item)
        if len(parser.rargs) >= 2:
            ImdbQueueManager.options['what'] = parser.rargs[1]
<commit_msg>imdb_queue now waits until feed_exit to remove movies from the queue, to make sure they have not been rejected or failed. fixes #503<commit_after>import logging
import datetime
from flexget.manager import Session
from flexget.plugin import register_plugin, PluginError, get_plugin_by_name, priority, register_parser_option
from flexget.manager import Base
from sqlalchemy import Column, Integer, String, DateTime, Boolean, Unicode
from flexget.utils.imdb import extract_id, ImdbSearch, ImdbParser, log as imdb_log
from flexget.utils.tools import str_to_boolean

log = logging.getLogger('imdb_queue')


class ImdbQueue(Base):

    __tablename__ = 'imdb_queue'

    id = Column(Integer, primary_key=True)
    imdb_id = Column(String)
    quality = Column(String)
    title = Column(Unicode)
    immortal = Column(Boolean)
    added = Column(DateTime)

    def __init__(self, imdb_id, quality, immortal):
        self.imdb_id = imdb_id
        self.quality = quality
        self.immortal = immortal
        self.added = datetime.datetime.now()

    def __str__(self):
        return '<ImdbQueue(imdb_id=%s,quality=%s,force=%s)>' % (self.imdb_id, self.quality, self.immortal)


class FilterImdbQueue(object):
    """
    Allows a queue of upcoming movies that will be forcibly allowed if the given quality matches

    Example:

    imdb_queue: yes
    """

    # Dict of entries accepted by this plugin {imdb_id: entry} format
    accepted_entries = {}

    def validator(self):
        from flexget import validator
        return validator.factory('boolean')

    @priority(129)
    def on_feed_filter(self, feed):
        for entry in feed.entries:
            # make sure the entry has IMDB fields filled
            try:
                get_plugin_by_name('imdb_lookup').instance.lookup(feed, entry)
            except PluginError:
                # no IMDB data, can't do anything
                continue

            # entry has imdb url
            if 'imdb_url' in entry:
                # get an imdb id
                if 'imdb_id' in entry and entry['imdb_id'] is not None:
                    imdb_id = entry['imdb_id']
                else:
                    imdb_id = extract_id(entry['imdb_url'])

                if not imdb_id:
                    log.warning("No imdb id could be determined for %s" % entry['title'])
                    continue

                if not 'quality' in entry:
                    log.warning('No quality found for %s, assigning unknown.' % entry['title'])
                    entry['quality'] = 'unknown'

                item = feed.session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).\
                                                     filter((ImdbQueue.quality == entry['quality']) | (ImdbQueue.quality == "ANY")).first()

                if item:
                    entry['immortal'] = item.immortal
                    log.info("Accepting %s from queue with quality %s. Force: %s" % (entry['title'], entry['quality'], entry['immortal']))
                    feed.accept(entry, 'imdb-queue - force: %s' % entry['immortal'])
                    # Keep track of entries we accepted, so they can be removed from queue on feed_exit if successful
                    self.accepted_entries[imdb_id] = entry
                else:
                    log.log(5, "%s not in queue with wanted quality, skipping" % entry['title'])

    def on_feed_exit(self, feed):
        """
        Removes any entries that have not been rejected by another plugin or failed from the queue.
        """
        for imdb_id, entry in self.accepted_entries.iteritems():
            if entry in feed.accepted and entry not in feed.failed:
                # If entry was not rejected or failed, remove from database
                item = feed.session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).first()
                feed.session.delete(item)
                log.debug('%s was successful, removing from imdb-queue' % entry['title'])


class ImdbQueueManager(object):
    """
    Handle IMDb queue management; add, delete and list
    """

    valid_actions = ['add', 'del', 'list']

    options = {}

    @staticmethod
    def optik_imdb_queue(option, opt, value, parser):
        """
        Callback for Optik
        --imdb-queue (add|del|list) [IMDB_URL|NAME] [quality]
        """
        if len(parser.rargs) == 0:
            print 'Usage: --imdb-queue (add|del|list) [IMDB_URL|NAME] [QUALITY] [FORCE]'
            # set some usage option so that feeds will be disabled later
            ImdbQueueManager.options['usage'] = True
            return

        ImdbQueueManager.options['action'] = parser.rargs[0].lower()

        if len(parser.rargs) == 1:
            return
        # 2 args is the minimum allowed (operation + item)
        if len(parser.rargs) >= 2:
            ImdbQueueManager.options['what'] = parser.rargs[1]

{'fraction_non_alphanumeric': 0.10177887863093897, 'fraction_numerical': 0.0025895068678225626, 'mean_word_length': 2.179312813171081, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4106596', 'n_tokens_mistral': 2601, 'n_tokens_neox': 2433, 'n_words': 715}