flexget.plugins.filter_imdb_queue
Covered: 153 lines
Missed: 155 lines
Skipped 103 lines
Percent: 49 %
  1
import logging
  2
import datetime
  3
from flexget.manager import Session
  4
from flexget.plugin import register_plugin, PluginError, get_plugin_by_name, priority, register_parser_option, register_feed_phase
  5
from flexget.manager import Base
  6
from sqlalchemy import Column, Integer, String, DateTime, Boolean, Unicode
  7
from flexget.utils.imdb import extract_id, ImdbSearch, ImdbParser, log as imdb_log
  8
from flexget.utils.tools import str_to_boolean
  9
from flexget.utils import qualities
 11
log = logging.getLogger('imdb_queue')
 14
class QueueError(Exception):
 15
    """Exception raised if there is an error with a queue operation"""
 17
    def __init__(self, message, errno=0):
 18
        self.message = message
 19
        self.errno = errno
 22
class ImdbQueue(Base):
 24
    __tablename__ = 'imdb_queue'
 26
    id = Column(Integer, primary_key=True)
 27
    imdb_id = Column(String)
 28
    quality = Column(String)
 29
    title = Column(Unicode)
 30
    immortal = Column(Boolean)
 31
    added = Column(DateTime)
 33
    def __init__(self, imdb_id, quality, immortal):
 34
        self.imdb_id = imdb_id
 35
        self.quality = quality
 36
        self.immortal = immortal
 37
        self.added = datetime.datetime.now()
 39
    def __str__(self):
 40
        return '<ImdbQueue(imdb_id=%s,quality=%s,force=%s)>' % (self.imdb_id, self.quality, self.immortal)
 43
class FilterImdbQueue(object):
 44
    """
 45
    Allows a queue of upcoming movies that will be forcibly allowed if the given quality matches
 47
    Example:
 49
    imdb_queue: yes
 50
    """
 53
    accepted_entries = {}
 55
    def validator(self):
 56
        from flexget import validator
 57
        return validator.factory('boolean')
 59
    @priority(129)
 60
    def on_feed_filter(self, feed):
 65
        rejected = []
 66
        for entry in feed.entries:
 68
            try:
 69
                get_plugin_by_name('imdb_lookup').instance.lookup(feed, entry)
 70
            except PluginError:
 72
                continue
 75
            if 'imdb_url' in entry:
 77
                if 'imdb_id' in entry and entry['imdb_id'] is not None:
 78
                    imdb_id = entry['imdb_id']
 79
                else:
 80
                    imdb_id = extract_id(entry['imdb_url'])
 82
                if not imdb_id:
 83
                    log.warning("No imdb id could be determined for %s" % entry['title'])
 84
                    continue
 86
                item = feed.session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).first()
 88
                if item:
 89
                    entry['immortal'] = item.immortal
 90
                    log.debug("Pre-accepting %s from queue" % (entry['title']))
 91
                    feed.accept(entry, 'imdb-queue pre-accept')
 93
                    self.accepted_entries[imdb_id] = entry
 94
                else:
 95
                    log.debugall("%s not in queue, skipping" % entry['title'])
 97
    def on_feed_imdbqueue(self, feed):
 98
        rejected = []
 99
        for entry in feed.entries:
100
            if entry['url'] == '':
101
                feed.reject(entry, 'imdb-queue - no URL in entry: '
102
                            '%s' % entry['title'])
103
                rejected.append(entry['title'])
104
                continue
106
            try:
107
                get_plugin_by_name('imdb_lookup').instance.lookup(feed, entry)
108
            except PluginError:
110
                continue
113
            if 'imdb_url' in entry:
115
                if 'imdb_id' in entry and entry['imdb_id'] is not None:
116
                    imdb_id = entry['imdb_id']
117
                else:
118
                    imdb_id = extract_id(entry['imdb_url'])
120
                if not imdb_id:
121
                    log.warning("No imdb id could be determined for %s" % entry['title'])
122
                    continue
124
                if not 'quality' in entry:
125
                    log.warning('No quality found for %s, assigning unknown.' % entry['title'])
126
                    entry['quality'] = 'unknown'
128
                item = feed.session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).first()
130
                if item:
132
                    minquality = qualities.parse_quality(item.quality)
133
                    if 'quality' in entry:
134
                        entry_quality = qualities.parse_quality(entry['quality'])
135
                    if entry_quality >= minquality:
136
                        entry['immortal'] = item.immortal
137
                        log.debug("found %s quality for %s. Need minimum %s" %
138
                                  (entry['title'], entry['quality'],
139
                                   item.quality))
140
                        log.info("Accepting %s from queue with quality %s. Force: %s" % (entry['title'], entry['quality'], entry['immortal']))
141
                        feed.accept(entry, 'imdb-queue - force: %s' % entry['immortal'])
143
                        self.accepted_entries[imdb_id] = entry
144
                    else:
145
                        log.debug("imdb-queue rejecting - found "
146
                                  "%s quality for %s. Need minimum %s" %
147
                                  (entry['title'], entry['quality'],
148
                                   item.quality))
151
                        entry['immortal'] = False
152
                        feed.reject(entry, 'imdb-queue quality '
153
                                    '%s below minimum %s for %s' %
154
                                    (entry_quality.name, minquality.name,
155
                                     entry['title']))
156
                else:
157
                    log.debugall("%s not in queue with wanted quality, skipping" % entry['title'])
158
        if len(rejected):
159
            log.info("Rejected due to no URL in entry (URLRewrite probably failed):"
160
                     " %s" % ', '.join(rejected))
162
    def on_feed_exit(self, feed):
163
        """
164
        Removes any entries that have not been rejected by another plugin or failed from the queue.
165
        """
166
        for imdb_id, entry in self.accepted_entries.iteritems():
167
            if entry in feed.accepted and entry not in feed.failed:
169
                item = feed.session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).first()
170
                feed.session.delete(item)
171
                log.debug('%s was successful, removing from imdb-queue' % entry['title'])
174
class ImdbQueueManager(object):
175
    """
176
    Handle IMDb queue management; add, delete and list
177
    """
179
    valid_actions = ['add', 'del', 'list']
181
    options = {}
183
    @staticmethod
184
    def optik_imdb_queue(option, opt, value, parser):
185
        """
186
        Callback for Optik
187
        --imdb-queue (add|del|list) [IMDB_URL|NAME] [quality]
188
        """
189
        if not parser.rargs:
190
            print 'Usage: --imdb-queue (add|del|list) [IMDB_URL|NAME] [QUALITY] [FORCE]'
192
            ImdbQueueManager.options['usage'] = True
193
            return
195
        ImdbQueueManager.options['action'] = parser.rargs[0].lower()
197
        if len(parser.rargs) == 1:
198
            return
200
        if len(parser.rargs) >= 2:
201
            ImdbQueueManager.options['what'] = parser.rargs[1]
204
        if len(parser.rargs) >= 3:
205
            ImdbQueueManager.options['quality'] = parser.rargs[2]
206
        else:
207
            ImdbQueueManager.options['quality'] = 'ANY' # TODO: Get default from config somehow?
210
        if len(parser.rargs) >= 4:
211
            ImdbQueueManager.options['force'] = str_to_boolean(parser.rargs[3])
212
        else:
213
            ImdbQueueManager.options['force'] = True
215
    def parse_what(self, what):
216
        """Given an imdb id or movie title, looks up from imdb and returns a dict with imdb_id and title keys"""
217
        imdb_id = extract_id(what)
218
        title = what
220
        if imdb_id:
222
            parser = ImdbParser()
223
            try:
224
                parser.parse('http://www.imdb.com/title/%s' % imdb_id)
225
            except Exception:
226
                raise QueueError('Error parsing info from imdb for %s' % imdb_id)
227
            if parser.name:
228
                title = parser.name
229
        else:
231
            print 'Searching imdb for %s' % what
232
            search = ImdbSearch()
233
            result = search.smart_match(what)
234
            if not result:
235
                raise QueueError('ERROR: Unable to find any such movie from imdb, use imdb url instead.')
236
            imdb_id = extract_id(result['url'])
237
            title = result['name']
239
        self.options['imdb_id'] = imdb_id
240
        self.options['title'] = title
241
        return {'title': title, 'imdb_id': imdb_id}
243
    def on_process_start(self, feed):
244
        """
245
        Handle IMDb queue management
246
        """
248
        if not self.options:
249
            return
251
        feed.manager.disable_feeds()
253
        if 'usage' in self.options:
254
            return
256
        action = self.options['action']
257
        if action not in self.valid_actions:
258
            self.error('Invalid action, valid actions are: ' + ', '.join(self.valid_actions))
259
            return
262
        if action != 'list':
263
            if not self.options.get('what'):
264
                self.error('No URL or NAME given')
265
                return
266
            else:
268
                try:
269
                    what = self.parse_what(self.options['what'])
270
                except QueueError, e:
271
                    print e.message
272
                else:
273
                    self.options.update(what)
275
            if not self.options.get('title') or not self.options.get('imdb_id'):
276
                print 'could not determine movie to add' # TODO: Rethink errors
277
                return
279
        from sqlalchemy.exceptions import OperationalError
280
        try:
281
            if action == 'add':
282
                try:
283
                    added = self.queue_add(title=self.options['title'], imdb_id=self.options['imdb_id'],
284
                               quality=self.options['quality'], force=self.options['force'])
285
                except QueueError, e:
286
                    print e.message
287
                    if e.errno == 1:
289
                        print 'Recognized qualities are %s' % ', '.join([qual.name for qual in qualities.all()])
290
                        print 'ANY is the default and can also be used explicitly to specify that quality should be ignored.'
291
                else:
292
                    print 'Added %s to queue with quality %s' % (added['title'], added['quality'])
293
            elif action == 'del':
294
                try:
295
                    title = self.queue_del(self.options['imdb_id'])
296
                except QueueError, e:
297
                    print e.message
298
                else:
299
                    print '%s removed from queue.' % title
300
            elif action == 'list':
301
                self.queue_list()
302
        except OperationalError:
303
            log.critical('OperationalError')
305
    def error(self, msg):
306
        print 'IMDb Queue error: %s' % msg
308
    def validate_quality(self, quality):
311
        if quality.upper() == 'ANY':
312
            return 'ANY'
313
        elif qualities.get(quality, False):
314
            return qualities.common_name(quality)
315
        else:
316
            raise QueueError('ERROR! Unknown quality `%s`' % quality, errno=1)
318
    def queue_add(self, title=None, imdb_id=None, quality='ANY', force=True):
319
        """Add an item to the queue with the specified quality"""
321
        if not title or not imdb_id:
323
            result = self.parse_what(imdb_id or title)
324
            title = result['title']
325
            imdb_id = result['imdb_id']
326
        quality = self.validate_quality(quality)
328
        session = Session()
331
        item = session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).first()
332
        if not item:
333
            item = ImdbQueue(imdb_id, quality, force)
334
            item.title = title
335
            session.add(item)
336
            session.commit()
337
            return {'title': title, 'imdb_id': imdb_id, 'quality': quality, 'force': force}
338
        else:
339
            raise QueueError('ERROR: %s is already in the queue' % title)
341
    def queue_del(self, imdb_id):
342
        """Delete the given item from the queue"""
344
        session = Session()
346
        item = session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).first()
347
        if item:
348
            title = item.title
349
            session.delete(item)
350
            session.commit()
351
            return title
352
        else:
353
            raise QueueError('%s is not in the queue' % imdb_id)
355
    def queue_edit(self, imdb_id, quality):
356
        """Change the required quality for a movie in the queue"""
357
        self.validate_quality(quality)
358
        session = Session()
360
        item = session.query(ImdbQueue).filter(ImdbQueue.imdb_id == imdb_id).first()
361
        if item:
362
            item.quality = quality
363
            session.commit()
364
            return item.title
365
        else:
366
            raise QueueError('%s is not in the queue' % imdb_id)
368
    def queue_list(self):
369
        """List IMDb queue"""
371
        items = self.queue_get()
372
        print '-' * 79
373
        print '%-10s %-45s %-8s %s' % ('IMDB id', 'Title', 'Quality', 'Force')
374
        print '-' * 79
375
        for item in items:
376
            print '%-10s %-45s %-8s %s' % (item.imdb_id, item.title, item.quality, item.immortal)
378
        if not items:
379
            print 'IMDB queue is empty'
381
        print '-' * 79
383
    def queue_get(self):
384
        """Get the current IMDb queue.
386
        Returns:
388
        List of ImdbQueue objects (detached from session)
389
        """
390
        session = Session()
391
        try:
392
            items = session.query(ImdbQueue).all()
393
            for item in items:
394
                if not item.title:
396
                    try:
397
                        item.title = self.parse_what(item.imdb_id)['title']
398
                    except QueueError:
399
                        item.title = 'N/A'
400
            return items
401
        finally:
402
            session.close()
404
register_plugin(FilterImdbQueue, 'imdb_queue')
405
register_plugin(ImdbQueueManager, 'imdb_queue_manager', builtin=True)
407
register_feed_phase(FilterImdbQueue, 'imdbqueue', after='urlrewrite')
409
register_parser_option('--imdb-queue', action='callback', callback=ImdbQueueManager.optik_imdb_queue,
410
                       help='(add|del|list) [IMDB_URL|NAME] [QUALITY]')