Rewrite load_imdb_info

* Use dict literal to create imdb_info * Use .get() for imdb_info in displayShow.mako and snatchSelection.mako
pymedusa · Apr 16, 2016 · abef7e0 · abef7e0
1 parent f7ea535
commit abef7e0
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 90 deletions.
diff --git a/gui/slick/views/displayShow.mako b/gui/slick/views/displayShow.mako
@@ -128,9 +128,7 @@
     % if show.imdb_info.get('year'):
                     (${show.imdb_info['year']}) -
     % endif
-    % if show.imdb_info.get('runtimes'):
-                    ${show.imdb_info['runtimes']} minutes
-    % endif
+                    ${show.imdb_info.get('runtimes') or show.runtime} minutes
                 </span>
 
                 <a href="${anon_url('http://www.imdb.com/title/', show.imdbid)}" rel="noreferrer" onclick="window.open(this.href, '_blank'); return false;" title="http://www.imdb.com/title/${show.imdbid}">
@@ -151,14 +149,14 @@
 
             <div id="tags">
                 <ul class="tags">
-                    % if ('genres' not in show.imdb_info or not show.imdb_info['genres']) and show.genre:
-                        % for genre in show.genre[1:-1].split('|'):
-                            <a href="${anon_url('http://trakt.tv/shows/popular/?genres=', genre.lower())}" target="_blank" title="View other popular ${genre} shows on trakt.tv."><li>${genre}</li></a>
-                        % endfor
-                    % elif 'genres' in show.imdb_info and show.imdb_info['genres']:
+                    % if show.imdb_info.get('genres'):
                         % for imdbgenre in show.imdb_info['genres'].replace('Sci-Fi','Science-Fiction').split('|'):
                             <a href="${anon_url('http://www.imdb.com/search/title?count=100&title_type=tv_series&genres=', imdbgenre.lower())}" target="_blank" title="View other popular ${imdbgenre} shows on IMDB."><li>${imdbgenre}</li></a>
                         % endfor
+                    % elif show.genre:
+                        % for genre in show.genre[1:-1].split('|'):
+                            <a href="${anon_url('http://trakt.tv/shows/popular/?genres=', genre.lower())}" target="_blank" title="View other popular ${genre} shows on trakt.tv."><li>${genre}</li></a>
+                        % endfor
                     % endif
                 </ul>
             </div>

diff --git a/gui/slick/views/snatchSelection.mako b/gui/slick/views/snatchSelection.mako
@@ -64,9 +64,7 @@
     % if show.imdb_info.get('year'):
                     (${show.imdb_info['year']}) -
     % endif
-    % if show.imdb_info.get('runtimes'):
-                    ${show.imdb_info['runtimes']} minutes
-    % endif
+                    ${show.imdb_info.get('runtimes') or show.runtime} minutes
                 </span>
 
                 <a href="${anon_url('http://www.imdb.com/title/', show.imdbid)}" rel="noreferrer" onclick="window.open(this.href, '_blank'); return false;" title="http://www.imdb.com/title/${show.imdbid}">
@@ -87,14 +85,14 @@
 
             <div id="tags">
                 <ul class="tags">
-                    % if ('genres' not in show.imdb_info or not show.imdb_info['genres']) and show.genre:
-                        % for genre in show.genre[1:-1].split('|'):
-                            <a href="${anon_url('http://trakt.tv/shows/popular/?genres=', genre.lower())}" target="_blank" title="View other popular ${genre} shows on trakt.tv."><li>${genre}</li></a>
-                        % endfor
-                    % elif 'genres' in show.imdb_info and show.imdb_info['genres']:
+                    % if show.imdb_info.get('genres'):
                         % for imdbgenre in show.imdb_info['genres'].replace('Sci-Fi','Science-Fiction').split('|'):
                             <a href="${anon_url('http://www.imdb.com/search/title?count=100&title_type=tv_series&genres=', imdbgenre.lower())}" target="_blank" title="View other popular ${imdbgenre} shows on IMDB."><li>${imdbgenre}</li></a>
                         % endfor
+                    % elif show.genre:
+                        % for genre in show.genre[1:-1].split('|'):
+                            <a href="${anon_url('http://trakt.tv/shows/popular/?genres=', genre.lower())}" target="_blank" title="View other popular ${genre} shows on trakt.tv."><li>${genre}</li></a>
+                        % endfor
                     % endif
                 </ul>
             </div>

diff --git a/sickbeard/show_queue.py b/sickbeard/show_queue.py
@@ -464,9 +464,9 @@ def run(self):
 
         logger.log(u"Retrieving show info from IMDb", logger.DEBUG)
         try:
-            self.show.loadIMDbInfo()
+            self.show.load_imdb_info()
         except imdb_exceptions.IMDbError as e:
-            logger.log(u" Something wrong on IMDb api: " + ex(e), logger.WARNING)
+            logger.log(u"Something wrong on IMDb api: " + ex(e), logger.WARNING)
         except Exception as e:
             logger.log(u"Error loading IMDb info: " + ex(e), logger.ERROR)
 
@@ -649,9 +649,9 @@ def run(self):
 
         logger.log(u"Retrieving show info from IMDb", logger.DEBUG)
         try:
-            self.show.loadIMDbInfo()
+            self.show.load_imdb_info()
         except imdb_exceptions.IMDbError as e:
-            logger.log(u" Something wrong on IMDb api: " + ex(e), logger.WARNING)
+            logger.log(u"Something wrong on IMDb api: " + ex(e), logger.WARNING)
         except Exception as e:
             logger.log(u"Error loading IMDb info: " + ex(e), logger.ERROR)
             logger.log(traceback.format_exc(), logger.DEBUG)

diff --git a/sickbeard/tv.py b/sickbeard/tv.py
@@ -487,7 +487,7 @@ def loadEpisodesFromDB(self):  # pylint: disable=too-many-locals
             curSeason = int(curResult["season"])
             curEpisode = int(curResult["episode"])
             curShowid = int(curResult['showid'])
-            curShowName = str(curResult['show_name'])
+            curShowName = unicode(curResult['show_name'])
 
             logger.log(u"%s: Loading %s episodes from DB" % (curShowid, curShowName), logger.DEBUG)
             deleteEp = False
@@ -877,84 +877,51 @@ def loadFromIndexer(self, cache=True, tvapi=None):
 
         self.status = getattr(myEp, 'status', 'Unknown')
 
-    def loadIMDbInfo(self):  # pylint: disable=too-many-branches
-
-        imdb_info = {
-            'imdb_id': self.imdbid,
-            'title': '',
-            'year': '',
-            'akas': [],
-            'runtimes': '',
-            'genres': [],
-            'countries': '',
-            'country_codes': [],
-            'certificates': [],
-            'rating': '',
-            'votes': '',
-            'last_update': ''
-        }
-
-        i = imdb.IMDb()
+    def load_imdb_info(self):  # pylint: disable=too-many-branches
+        """Load all required show information from IMDb with IMDbPY."""
+        imdb_api = imdb.IMDb()
         if not self.imdbid:
-            self.imdbid = i.title2imdbID(self.name, kind='tv series')
+            self.imdbid = imdb_api.title2imdbID(self.name, kind='tv series')
 
         if not self.imdbid:
-            logger.log(str(self.indexerid) + u": Not loading show info from IMDb, because we don't know the imdbid", logger.DEBUG)
+            logger.log(u"{0}: Not loading show info from IMDb, "
+                       u"because we don't know its ID".format(self.indexerid), logger.ERROR)
             return
 
-        logger.log(str(self.indexerid) + u": Loading show info from IMDb", logger.DEBUG)
-
-        imdbTv = i.get_movie(str(re.sub(r"[^0-9]", "", self.imdbid)))
-
-        for key in [x for x in imdb_info.keys() if x.replace('_', ' ') in imdbTv.keys()]:
-            # Store only the first value for string type
-            if isinstance(imdb_info[key], basestring) and isinstance(imdbTv.get(key.replace('_', ' ')), list):
-                imdb_info[key] = imdbTv.get(key.replace('_', ' '))[0]
-            else:
-                imdb_info[key] = imdbTv.get(key.replace('_', ' '))
-
-        # Filter only the value
-        if imdb_info['runtimes']:
-            imdb_info['runtimes'] = re.search(r'\d+', imdb_info['runtimes']).group(0)
-        else:
-            imdb_info['runtimes'] = self.runtime
-
-        if imdb_info['akas']:
-            imdb_info['akas'] = '|'.join(imdb_info['akas'])
-        else:
-            imdb_info['akas'] = ''
+        # Make sure we only use one ID
+        imdb_id = self.imdbid.split(',')[0]
+
+        logger.log(u'{0}: Loading show info from IMDb with ID: {1}'.format(
+                   self.indexerid, imdb_id), logger.DEBUG)
+
+        # Remove first two chars from ID
+        imdb_obj = imdb_api.get_movie(imdb_id[2:])
+
+        self.imdb_info = {
+            'imdb_id': imdb_id,
+            'title': imdb_obj.get('title', ''),
+            'year': imdb_obj.get('year', ''),
+            'akas': '|'.join(imdb_obj.get('akas', '')),
+            'genres': '|'.join(imdb_obj.get('genres', '')),
+            'countries': '|'.join(imdb_obj.get('countries', '')),
+            'country_codes': '|'.join(imdb_obj.get('country codes', '')),
+            'rating': imdb_obj.get('rating', ''),
+            'votes': imdb_obj.get('votes', ''),
+            'last_update': datetime.date.today().toordinal()
+        }
 
-        # Join all genres in a string
-        if imdb_info['genres']:
-            imdb_info['genres'] = '|'.join(imdb_info['genres'])
-        else:
-            imdb_info['genres'] = ''
+        if imdb_obj.get('runtimes'):
+            self.imdb_info['runtimes'] = re.search(r'\d+', imdb_obj['runtimes'][0]).group(0)
 
         # Get only the production country certificate if any
-        if imdb_info['certificates'] and imdb_info['countries']:
-            dct = {}
-            try:
-                for item in imdb_info['certificates']:
-                    dct[item.split(':')[0]] = item.split(':')[1]
-
-                imdb_info['certificates'] = dct[imdb_info['countries']]
-            except Exception:
-                imdb_info['certificates'] = ''
-
-        else:
-            imdb_info['certificates'] = ''
-
-        if imdb_info['country_codes']:
-            imdb_info['country_codes'] = '|'.join(imdb_info['country_codes'])
-        else:
-            imdb_info['country_codes'] = ''
-
-        imdb_info['last_update'] = datetime.date.today().toordinal()
+        if imdb_obj.get('certificates') and imdb_obj.get('countries'):
+            for certificate in imdb_obj['certificates']:
+                if certificate.split(':')[0] in imdb_obj['countries']:
+                    self.imdb_info['certificates'] = certificate.split(':')[1]
+                    break
 
-        # Rename dict keys without spaces for DB upsert
-        self.imdb_info = dict(
-            (k.replace(' ', '_'), k(v) if hasattr(v, 'keys') else v) for k, v in imdb_info.iteritems())
-        logger.log(str(self.indexerid) + u": Obtained info from IMDb ->" + str(self.imdb_info), logger.DEBUG)
+        logger.log(u'{0}: Obtained info from IMDb: {1}'.format(
+                   self.indexerid, self.imdb_info), logger.DEBUG)
 
     def nextEpisode(self):
         logger.log(str(self.indexerid) + ": Finding the episode which airs next", logger.DEBUG)