Merge pull request #12 from zeroquinc/trakt

enhancement: optimize trakt.py performance and code logic

config.py

@@ -27,7 +27,7 @@
 
 IMDB_BUTTON_ENABLED = True  # Set this to False to disable the IMDb button
 TMDB_BUTTON_ENABLED = False  # Set this to False to disable the TMDB button
-TRAKT_BUTTON_ENABLED = False  # Set this to False to disable the Trakt buttons
+TRAKT_BUTTON_ENABLED = True  # Set this to False to disable the Trakt buttons
 LETTERBOXD_BUTTON_ENABLED = False  # Set this to False to disable the Letterboxd button
 
 """

src/trakt.py

@@ -1,73 +1,46 @@
 import requests
-from urllib.parse import quote
-
 from .custom_logger import logger
-from .globals import TRAKT_API_KEY, TMDB_API_KEY
-
-"""
-This file contains functions to fetch data from the Trakt API.
-"""
+from .globals import TRAKT_API_KEY
 
-# Function to get the TMDB ID of a media
 def get_tmdb_id_trakt(info, media_type):
-    # If the media type is a channel, return None
     if media_type == 'channel':
         return None
 
-    tmdb_id = None
-    if info['type'] == 'episode':
-        tmdb_id = get_tmdb_id_for_media(info, 'tv')
-    elif info['type'] == 'movie':
-        tmdb_id = get_tmdb_id_for_media(info, 'movie')
-
+    tmdb_id = get_tmdb_id_for_media(info) if info['type'] in ('episode', 'movie') else None
     logger.debug(f"TMDB ID: {tmdb_id}")
     return tmdb_id
 
-# Function to fetch the tmdb_id of the media
-def get_tmdb_id_for_media(info, media_type):
-    if 'uniqueid' in info and 'tmdb' in info['uniqueid']:
-        logger.debug("Found uniqueid in info")
-        return info['uniqueid']['tmdb']
-    else:
-        logger.debug("Can't find TMDB ID in uniqueid")
-        return None
+def get_tmdb_id_for_media(info):
+    tmdb_id = info.get('uniqueid', {}).get('tmdb')
+    logger.debug(f"{'Found' if tmdb_id else 'Can\'t find'} TMDB ID in uniqueid")
+    return tmdb_id
 
-# Function to get the Trakt URL of a media
 def get_trakt_url(tmdb_id, media_type):
-        # Return None if tmdb_id is None
     if tmdb_id is None:
         return None
-
-    # Define the base URL
-    base_url = "https://trakt.tv/"
 
-    # Check the media type
-    if media_type == "movie":
-        media_url = base_url + "movies/"
-        search_type = "movie"
-    elif media_type == "tv":
-        media_url = base_url + "shows/"
-        search_type = "show"
-    else:
+    media_urls = {"movie": "movies/", "tv": "shows/"}
+    search_types = {"movie": "movie", "tv": "show"}
+
+    media_url = f"https://trakt.tv/{media_urls.get(media_type)}"
+    search_type = search_types.get(media_type)
+
+    if not media_url or not search_type:
         return None
 
     search_url = f"https://api.trakt.tv/search/tmdb/{tmdb_id}?type={search_type}"
-
-    # Send a GET request to the search endpoint
     headers = {"trakt-api-key": TRAKT_API_KEY}
     response = requests.get(search_url, headers=headers)
 
-    # Check if the request was successful
     if response.status_code == 200:
-        # Parse the JSON response
-        data = response.json()
-        logger.debug(f"Trakt search response: {data}")
-
-        trakt_slug = data[0][search_type]["ids"]["slug"]
-        trakt_url = media_url + trakt_slug
-
-        logger.debug(f"Generated Trakt URL: {trakt_url}")
-        return trakt_url
-    else:
-        logger.error(f"Failed to get Trakt URL: {response.content}")
-        return None
+        return trakt_url(response, search_type, media_url)
+    logger.error(f"Failed to get Trakt URL: {response.content}")
+    return None
+
+def trakt_url(response, search_type, media_url):
+    data = response.json()
+    logger.debug(f"Trakt search response: {data}")
+    trakt_slug = data[0][search_type]["ids"]["slug"]
+    trakt_url = media_url + trakt_slug
+    logger.debug(f"Generated Trakt URL: {trakt_url}")
+    return trakt_url