diff --git a/couchpotato/core/providers/automation/imdb/__init__.py b/couchpotato/core/providers/automation/imdb/__init__.py
index ee804af1..546cba97 100644
--- a/couchpotato/core/providers/automation/imdb/__init__.py
+++ b/couchpotato/core/providers/automation/imdb/__init__.py
@@ -38,21 +38,23 @@ config = [{
'description': 'Import movies from IMDB Charts',
'options': [
{
- 'name': 'automation_enabled',
+ 'name': 'automation_providers_enabled',
'default': False,
'type': 'enabler',
},
{
- 'name': 'automation_charts_theaters_use',
- 'type': 'checkbox',
+ 'name': 'automation_charts_theater',
+ 'type': 'bool',
'label': 'In Theaters',
'description': 'New Movies In-Theaters chart',
+ 'default': True,
},
{
- 'name': 'automation_charts_top250_use',
- 'type': 'checkbox',
+ 'name': 'automation_charts_top250',
+ 'type': 'bool',
'label': 'TOP 250',
'description': 'IMDB TOP 250 chart',
+ 'default': True,
},
],
},
diff --git a/couchpotato/core/providers/automation/imdb/main.py b/couchpotato/core/providers/automation/imdb/main.py
index 0d494949..c4aef7f1 100644
--- a/couchpotato/core/providers/automation/imdb/main.py
+++ b/couchpotato/core/providers/automation/imdb/main.py
@@ -1,90 +1,41 @@
+import traceback
+
from bs4 import BeautifulSoup
+from couchpotato import fireEvent
from couchpotato.core.helpers.rss import RSS
from couchpotato.core.helpers.variable import getImdb, splitString, tryInt
+
from couchpotato.core.logger import CPLog
from couchpotato.core.providers.automation.base import Automation
-import re
-import traceback
+
+from couchpotato.core.providers.base import MultiProvider
+
log = CPLog(__name__)
-class IMDB(Automation, RSS):
+class IMDB(MultiProvider):
+
+ def getTypes(self):
+ return [IMDBWatchlist, IMDBAutomation]
+
+
+class IMDBBase(Automation, RSS):
interval = 1800
- chart_urls = {
- 'theater': 'http://www.imdb.com/movies-in-theaters/',
- 'top250': 'http://www.imdb.com/chart/top',
- }
+ def getInfo(self, imdb_id):
+ return fireEvent('movie.info', identifier = imdb_id, merge = True)
+class IMDBWatchlist(IMDBBase):
+
+ enabled_option = 'automation_enabled'
+
def getIMDBids(self):
movies = []
- # Handle Chart URLs
- if self.conf('automation_charts_theaters_use'):
- log.debug('Started IMDB chart: %s', self.chart_urls['theater'])
- data = self.getHTMLData(self.chart_urls['theater'])
- if data:
- html = BeautifulSoup(data)
-
- try:
- result_div = html.find('div', attrs = {'id': 'main'})
-
- entries = result_div.find_all('div', attrs = {'itemtype': 'http://schema.org/Movie'})
-
- for entry in entries:
- title = entry.find('h4', attrs = {'itemprop': 'name'}).getText()
-
- log.debug('Identified title: %s', title)
- result = re.search('(.*) \((.*)\)', title)
-
- if result:
- name = result.group(1)
- year = result.group(2)
-
- imdb = self.search(name, year)
-
- if imdb and self.isMinimalMovie(imdb):
- movies.append(imdb['imdb'])
-
- except:
- log.error('Failed loading IMDB chart results from %s: %s', (self.chart_urls['theater'], traceback.format_exc()))
-
- if self.conf('automation_charts_top250_use'):
- log.debug('Started IMDB chart: %s', self.chart_urls['top250'])
- data = self.getHTMLData(self.chart_urls['top250'])
- if data:
- html = BeautifulSoup(data)
-
- try:
- result_div = html.find('div', attrs = {'id': 'main'})
-
- result_table = result_div.find_all('table')[1]
- entries = result_table.find_all('tr')
-
- for entry in entries[1:]:
- title = entry.find_all('td')[2].getText()
-
- log.debug('Identified title: %s', title)
- result = re.search('(.*) \((.*)\)', title)
-
- if result:
- name = result.group(1)
- year = result.group(2)
-
- imdb = self.search(name, year)
-
- if imdb and self.isMinimalMovie(imdb):
- movies.append(imdb['imdb'])
-
- except:
- log.error('Failed loading IMDB chart results from %s: %s', (self.chart_urls['theater'], traceback.format_exc()))
-
-
- # Handle Watchlists
watchlist_enablers = [tryInt(x) for x in splitString(self.conf('automation_urls_use'))]
watchlist_urls = splitString(self.conf('automation_urls'))
@@ -103,9 +54,47 @@ class IMDB(Automation, RSS):
for imdb in imdbs:
movies.append(imdb)
+ if self.shuttingDown():
+ break
+
except:
log.error('Failed loading IMDB watchlist: %s %s', (url, traceback.format_exc()))
-
- # Return the combined resultset
+ return movies
+
+
+class IMDBAutomation(IMDBBase):
+
+ enabled_option = 'automation_providers_enabled'
+
+ chart_urls = {
+ 'theater': 'http://www.imdb.com/movies-in-theaters/',
+ 'top250': 'http://www.imdb.com/chart/top',
+ }
+
+ def getIMDBids(self):
+
+ movies = []
+
+ for url in self.chart_urls:
+ if self.conf('automation_charts_%s' % url):
+ data = self.getHTMLData(self.chart_urls[url])
+ if data:
+ html = BeautifulSoup(data)
+
+ try:
+ result_div = html.find('div', attrs = {'id': 'main'})
+ imdb_ids = getImdb(str(result_div), multiple = True)
+
+ for imdb_id in imdb_ids:
+ info = self.getInfo(imdb_id)
+ if info and self.isMinimalMovie(info):
+ movies.append(imdb_id)
+
+ if self.shuttingDown():
+ break
+
+ except:
+ log.error('Failed loading IMDB chart results from %s: %s', (url, traceback.format_exc()))
+
return movies