rewrite tmdb.py

movies/tmdb.py

@@ -1,5 +1,4 @@
 import asyncio
-import copy
 import logging
 import random
 
@@ -13,49 +12,52 @@ from movies.classes import Movie, TVShow, Season, Episode
 _BASE_API_URL = "https://api.themoviedb.org/3"
 _BASE_IMAGE_URL = None
 
+_TMDB_REQUEST_SEMAPHORE = asyncio.Semaphore(config.TMDB_CONCURRENT_REQUESTS_LIMIT)
+
+
+async def _make_tmdb_request(session: aiohttp.ClientSession, url: str, params: dict = None) -> dict:
+    async with _TMDB_REQUEST_SEMAPHORE:
+        headers = {
+            "Authorization": f"Bearer {config.TMDB_API_KEY}",
+            "Accept": "application/json"
+        }
+        request_params = {"language": config.TMDB_LANG, **(params or {})}
+
+        try:
+            async with session.get(url=url, headers=headers, params=request_params) as response:
+                response.raise_for_status()
+                return await response.json()
+        except aiohttp.ClientResponseError as e:
+            logging.error("TMDB request failed for %s: Status %s, Response: %s", url, e.status, e.message)
+            raise
+        except aiohttp.ClientError as e:
+            logging.error("Network error during TMDB request for %s: %s", url, e)
+            raise
+
+
+async def _get_client_session() -> aiohttp.ClientSession:
+    proxy = random.choice(config.PROXIES) if config.PROXIES else None
+    return aiohttp.ClientSession(proxy=proxy)
+
 
 @cached(TTLCache(maxsize=config.CACHE_MAXSIZE, ttl=config.CACHE_TTL))
-async def _fetch_tmdb_data():
+async def _fetch_tmdb_configuration():
     global _BASE_IMAGE_URL
 
     logging.info("Fetching TMDB configuration data")
 
-    headers = {
+    async with await _get_client_session() as session:
-        "Authorization": f"Bearer {config.TMDB_API_KEY}",
+        response_json = await _make_tmdb_request(session, f"{_BASE_API_URL}/configuration")
-        "Accept": "application/json"
-    }
-
-    proxy = random.choice(config.PROXIES) if config.PROXIES else None
-    async with aiohttp.ClientSession(proxy=proxy) as session:
-        async with session.get(
-                url=f"{_BASE_API_URL}/configuration",
-                headers=headers,
-                params={"language": config.TMDB_LANG}
-        ) as response:
-            assert response.status == 200, "Failed to fetch TMDB configuration data"
-
-            response_json = await response.json()
 
     _BASE_IMAGE_URL = response_json["images"]["base_url"]
 
 
 @cached(TTLCache(maxsize=config.CACHE_MAXSIZE, ttl=config.CACHE_TTL))
 async def _fetch_movie(movie: Movie) -> Movie:
-    headers = {
+    logging.info("Fetching movie data for TMDB ID: %s", movie.tmdb_id)
-        "Authorization": f"Bearer {config.TMDB_API_KEY}",
-        "Accept": "application/json"
-    }
 
-    proxy = random.choice(config.PROXIES) if config.PROXIES else None
+    async with await _get_client_session() as session:
-    async with aiohttp.ClientSession(proxy=proxy) as session:
+        response_json = await _make_tmdb_request(session, f"{_BASE_API_URL}/movie/{movie.tmdb_id}")
-        async with session.get(
-                url=f"{_BASE_API_URL}/movie/{movie.tmdb_id}",
-                headers=headers,
-                params={"language": config.TMDB_LANG}
-        ) as response:
-            assert response.status == 200, f"Failed to fetch information for {movie.tmdb_id}"
-
-            response_json = await response.json()
 
     return Movie(
         tmdb_id=movie.tmdb_id,
@@ -68,73 +70,58 @@ async def _fetch_movie(movie: Movie) -> Movie:
         adult=response_json.get("adult"),
         og_title=response_json.get("original_title"),
         homepage=response_json.get("homepage"),
-        poster_url=f"{_BASE_IMAGE_URL}original{response_json["poster_path"]}" if response_json.get(
+        poster_url=f"{_BASE_IMAGE_URL}original{response_json['poster_path']}" if response_json.get(
             "poster_path") else None,
-        backdrop_url=f"{_BASE_IMAGE_URL}original{response_json["backdrop_path"]}" if response_json.get(
+        backdrop_url=f"{_BASE_IMAGE_URL}original{response_json['backdrop_path']}" if response_json.get(
             "backdrop_path") else None,
         release_date=response_json.get("release_date"),
         genres=tuple(sorted([genre["name"].capitalize() for genre in response_json.get("genres", [])], key=len)),
     )
 
 
+async def _fetch_tv_season_data(session: aiohttp.ClientSession, tv_show_id: int, season_number: int) -> dict:
+    logging.debug("Fetching season %s data for TV show TMDB ID: %s", season_number, tv_show_id)
+    return await _make_tmdb_request(session, f"{_BASE_API_URL}/tv/{tv_show_id}/season/{season_number}")
+
+
 @cached(TTLCache(maxsize=config.CACHE_MAXSIZE, ttl=config.CACHE_TTL))
 async def _fetch_tv_show(tv_show: TVShow) -> TVShow:
-    headers = {
+    logging.info("Fetching TV show data for TMDB ID: %s", tv_show.tmdb_id)
-        "Authorization": f"Bearer {config.TMDB_API_KEY}",
-        "Accept": "application/json"
-    }
 
-    proxy = random.choice(config.PROXIES) if config.PROXIES else None
+    async with await _get_client_session() as session:
-    async with aiohttp.ClientSession(proxy=proxy) as session:
+        tv_response_json = await _make_tmdb_request(session, f"{_BASE_API_URL}/tv/{tv_show.tmdb_id}")
-        async with session.get(
-                url=f"{_BASE_API_URL}/tv/{tv_show.tmdb_id}",
-                headers=headers,
-                params={"language": config.TMDB_LANG}
-        ) as tv_response:
-            assert tv_response.status == 200, f"Failed to fetch information for {tv_show.tmdb_id}"
 
-            tv_response_json = await tv_response.json()
+        season_tasks = [
+            _fetch_tv_season_data(session, tv_show.tmdb_id, season.season_number)
+            for season in tv_show.seasons
+        ]
+        seasons_data_jsons = await asyncio.gather(*season_tasks)
 
-        seasons_responses_jsons = {}
+    raw_seasons_map = {s.season_number: s for s in tv_show.seasons}
-        for i in range(len(tv_show.seasons)):
-            season_number = tv_show.seasons[i].season_number
 
-            async with session.get(
+    processed_seasons_list = []
-                    url=f"{_BASE_API_URL}/tv/{tv_show.tmdb_id}/season/{season_number}",
+    for season_response_json in seasons_data_jsons:
-                    headers=headers,
+        season_number = season_response_json["season_number"]
-                    params={"language": config.TMDB_LANG}
+        raw_season = raw_seasons_map.get(season_number)
-            ) as season_response:
-                assert tv_response.status == 200, f"Failed to fetch information for {tv_show.tmdb_id}"
 
-                seasons_responses_jsons[season_number] = await season_response.json()
+        if not raw_season:
+            logging.warning("Skipping season %s for TV show %s as it's not in raw data.", season_number,
+                            tv_show.tmdb_id)
+            continue
+
+        raw_episodes_map = {e.episode_number: e for e in raw_season.episodes}
+        processed_episodes = []
 
-    seasons_dict = {}
-    for season_number, season_response_json in seasons_responses_jsons.items():
         for episode_response_json in season_response_json.get("episodes", []):
             episode_number = episode_response_json["episode_number"]
+            raw_episode = raw_episodes_map.get(episode_number)
 
-            raw_season = None
+            if not raw_episode:
-            for season in tv_show.seasons:
+                logging.warning("Skipping episode %s:%s for TV show %s as it's not in raw data.",
-                if season.season_number == season_number:
+                                season_number, episode_number, tv_show.tmdb_id)
-                    raw_season = season
-                    break
-
-            if raw_season is None:
                 continue
 
-            raw_episode = None
+            processed_episodes.append(Episode(
-            for episode in raw_season.episodes:
-                if episode.episode_number == episode_number:
-                    raw_episode = episode
-                    break
-
-            if raw_episode is None:
-                continue
-
-            if not seasons_dict.get(season_number):
-                seasons_dict[season_number] = []
-
-            episode = Episode(
                 episode_number=episode_number,
                 season_number=season_number,
                 runtime=episode_response_json.get("runtime"),
@@ -143,54 +130,43 @@ async def _fetch_tv_show(tv_show: TVShow) -> TVShow:
                 file_url=raw_episode.file_url,
                 episode_type=episode_response_json.get("episode_type"),
                 release_date=episode_response_json.get("air_date"),
-            )
+            ))
-            seasons_dict[season_number].append(episode)
 
-    seasons = []
+        processed_seasons_list.append(Season(
-    for season_number, season_response_json in seasons_responses_jsons.items():
-        raw_season = None
-        for season in tv_show.seasons:
-            if season.season_number == season_number:
-                raw_season = season
-                break
-
-        if raw_season is None:
-            continue
-
-        season = Season(
             season_number=season_number,
-            episodes_count=len(seasons_dict[season_number]),
+            episodes_count=len(processed_episodes),
             vote_average=season_response_json.get("vote_average"),
             title=season_response_json.get("name", str(season_number)),
-            poster_url=f"{_BASE_IMAGE_URL}original{season_response_json["poster_path"]}" if season_response_json.get(
+            poster_url=f"{_BASE_IMAGE_URL}original{season_response_json['poster_path']}" if season_response_json.get(
                 "poster_path") else None,
             release_date=season_response_json.get("air_date"),
-            episodes=tuple(copy.deepcopy(seasons_dict[season_number]))
+            episodes=tuple(processed_episodes)
-        )
+        ))
-        seasons.append(season)
+
+    processed_seasons_list.sort(key=lambda s: s.season_number)
 
     return TVShow(
         tmdb_id=tv_show.tmdb_id,
         number_of_episodes=tv_show.number_of_episodes,
-        number_of_seasons=len(seasons),
+        number_of_seasons=len(processed_seasons_list),
         vote_average=tv_response_json.get("vote_average"),
         adult=tv_response_json.get("adult"),
         title=tv_response_json.get("name"),
         og_title=tv_response_json.get("original_name"),
         homepage=tv_response_json.get("homepage"),
-        poster_url=f"{_BASE_IMAGE_URL}original{tv_response_json["poster_path"]}" if tv_response_json.get(
+        poster_url=f"{_BASE_IMAGE_URL}original{tv_response_json['poster_path']}" if tv_response_json.get(
             "poster_path") else None,
-        backdrop_url=f"{_BASE_IMAGE_URL}original{tv_response_json["backdrop_path"]}" if tv_response_json.get(
+        backdrop_url=f"{_BASE_IMAGE_URL}original{tv_response_json['backdrop_path']}" if tv_response_json.get(
             "backdrop_path") else None,
         release_date=tv_response_json.get("first_air_date"),
-        in_production=tv_response_json.get("production"),
+        in_production=tv_response_json.get("in_production"),
-        seasons=tuple(seasons),
+        seasons=tuple(processed_seasons_list),
         genres=tuple(sorted([genre["name"].capitalize() for genre in tv_response_json.get("genres", [])], key=len)),
     )
 
 
 async def _fetch_data(content: Movie | TVShow) -> Movie | TVShow:
-    logging.info("Fetching data for %s", content.tmdb_id)
+    logging.info("Fetching data for %s: %s", content.type, content.tmdb_id)
 
     match content.type:
         case "movie":
@@ -204,7 +180,7 @@ async def _fetch_data(content: Movie | TVShow) -> Movie | TVShow:
 async def fetch_all_data(contents: list[Movie | TVShow]) -> list[Movie | TVShow]:
     logging.info("Fetching data for %s contents", len(contents))
 
-    await _fetch_tmdb_data()
+    await _fetch_tmdb_configuration()
 
     tasks = [_fetch_data(content) for content in contents]
     return await asyncio.gather(*tasks)