From 6b1daa61a0106d6557047059348ab6e7caecb693 Mon Sep 17 00:00:00 2001 From: Goldy <153996346+g0ldyy@users.noreply.github.com> Date: Thu, 4 Jul 2024 01:25:32 +0200 Subject: [PATCH] fix: 0.30ms faster :pro: --- comet/api/stream.py | 48 ++++++++++++++------------------ comet/utils/general.py | 62 ++++++++++++++++-------------------------- comet/utils/models.py | 8 +++--- 3 files changed, 48 insertions(+), 70 deletions(-) diff --git a/comet/api/stream.py b/comet/api/stream.py index 91b344e..94d1012 100644 --- a/comet/api/stream.py +++ b/comet/api/stream.py @@ -110,14 +110,14 @@ async def stream(request: Request, b64config: str, type: str, id: str): balanced_hashes = await get_balanced_hashes(sorted_ranked_files, config) results = [] - for hash in sorted_ranked_files: - for resolution in balanced_hashes: - if hash in balanced_hashes[resolution]: + for hash, hash_data in sorted_ranked_files.items(): + for resolution, hash_list in balanced_hashes.items(): + if hash in hash_list: results.append( { - "name": f"[RD⚡] Comet {sorted_ranked_files[hash]['data']['resolution'][0] if len(sorted_ranked_files[hash]['data']['resolution']) > 0 else 'Unknown'}", - "title": f"{sorted_ranked_files[hash]['data']['title']}\n💾 {bytes_to_size(sorted_ranked_files[hash]['data']['size'])}", - "url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{sorted_ranked_files[hash]['data']['index']}", + "name": f"[RD⚡] Comet {hash_data['data']['resolution'][0] if hash_data['data']['resolution'] else 'Unknown'}", + "title": f"{hash_data['data']['title']}\n💾 {bytes_to_size(hash_data['data']['size'])}", + "url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{hash_data['data']['index']}", } ) @@ -133,24 +133,18 @@ async def stream(request: Request, b64config: str, type: str, id: str): f"Start of {indexer_manager_type} search for {logName} with indexers {config['indexers']}" ) - tasks = [] - tasks.append( - get_indexer_manager(session, indexer_manager_type, config["indexers"], name) - ) + search_terms = [name] if type == "series": - tasks.append( - get_indexer_manager( - session, - indexer_manager_type, - config["indexers"], - f"{name} S0{season}E0{episode}", - ) - ) + search_terms.append(f"{name} S0{season}E0{episode}") + tasks = [ + get_indexer_manager(session, indexer_manager_type, config["indexers"], term) + for term in search_terms + ] search_response = await asyncio.gather(*tasks) torrents = [] for results in search_response: - if results == None: + if results is None: continue for result in results: @@ -191,9 +185,9 @@ async def stream(request: Request, b64config: str, type: str, id: str): logger.info( f"{zilean_hashes_count} torrents found for {logName} with Zilean API" ) - except: + except Exception as e: logger.warning( - f"Exception while getting torrents for {logName} with Zilean API" + f"Exception while getting torrents for {logName} with Zilean API: {e}" ) if len(torrents) == 0: @@ -310,14 +304,14 @@ async def stream(request: Request, b64config: str, type: str, id: str): balanced_hashes = await get_balanced_hashes(sorted_ranked_files, config) results = [] - for hash in sorted_ranked_files: - for resolution in balanced_hashes: - if hash in balanced_hashes[resolution]: + for hash, hash_data in sorted_ranked_files.items(): + for resolution, hash_list in balanced_hashes.items(): + if hash in hash_list: results.append( { - "name": f"[RD⚡] Comet {sorted_ranked_files[hash]['data']['resolution'][0] if len(sorted_ranked_files[hash]['data']['resolution']) > 0 else 'Unknown'}", - "title": f"{sorted_ranked_files[hash]['data']['title']}\n💾 {bytes_to_size(sorted_ranked_files[hash]['data']['size'])}", - "url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{sorted_ranked_files[hash]['data']['index']}", + "name": f"[RD⚡] Comet {hash_data['data']['resolution'][0] if hash_data['data']['resolution'] else 'Unknown'}", + "title": f"{hash_data['data']['title']}\n💾 {bytes_to_size(hash_data['data']['size'])}", + "url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{hash_data['data']['index']}", } ) diff --git a/comet/utils/general.py b/comet/utils/general.py index af701d2..f0b6502 100644 --- a/comet/utils/general.py +++ b/comet/utils/general.py @@ -263,7 +263,7 @@ async def get_indexer_manager( async def get_torrent_hash( session: aiohttp.ClientSession, indexer_manager_type: str, torrent: dict ): - if "InfoHash" in torrent and torrent["InfoHash"] != None: + if "InfoHash" in torrent and torrent["InfoHash"] is not None: return torrent["InfoHash"] if "infoHash" in torrent: @@ -302,41 +302,30 @@ async def get_torrent_hash( async def get_balanced_hashes(hashes: dict, config: dict): max_results = config["maxResults"] config_resolutions = config["resolutions"] - config_languages = config["languages"] + config_languages = {language.replace("_", " ").capitalize() for language in config["languages"]} + include_all_languages = "All" in config_languages + include_all_resolutions = "All" in config_resolutions + include_unknown_resolution = include_all_resolutions or "Unknown" in config_resolutions hashes_by_resolution = {} - for hash in hashes: - if ( - "All" not in config_languages - and not hashes[hash]["data"]["is_multi_audio"] - and not any( - language.replace("_", " ").capitalize() - in hashes[hash]["data"]["language"] - for language in config_languages - ) - ): + for hash, hash_data in hashes.items(): + hash_info = hash_data["data"] + if not include_all_languages and not hash_info["is_multi_audio"] and not any(lang in hash_info["language"] for lang in config_languages): continue - resolution = hashes[hash]["data"]["resolution"] - if len(resolution) == 0: - if "All" not in config_resolutions and "Unknown" not in config_resolutions: + resolution = hash_info["resolution"] + if not resolution: + if not include_unknown_resolution: continue - - if "Unknown" not in hashes_by_resolution: - hashes_by_resolution["Unknown"] = [hash] + resolution_key = "Unknown" + else: + resolution_key = resolution[0] + if not include_all_resolutions and resolution_key not in config_resolutions: continue - hashes_by_resolution["Unknown"].append(hash) - continue - - if "All" not in config_resolutions and resolution[0] not in config_resolutions: - continue - - if resolution[0] not in hashes_by_resolution: - hashes_by_resolution[resolution[0]] = [hash] - continue - - hashes_by_resolution[resolution[0]].append(hash) + if resolution_key not in hashes_by_resolution: + hashes_by_resolution[resolution_key] = [] + hashes_by_resolution[resolution_key].append(hash) if max_results == 0: return hashes_by_resolution @@ -346,25 +335,20 @@ async def get_balanced_hashes(hashes: dict, config: dict): extra_hashes = max_results % total_resolutions balanced_hashes = {} - for resolution, hashes in hashes_by_resolution.items(): - selected_count = hashes_per_resolution - + for resolution, hash_list in hashes_by_resolution.items(): + selected_count = hashes_per_resolution + (1 if extra_hashes > 0 else 0) + balanced_hashes[resolution] = hash_list[:selected_count] if extra_hashes > 0: - selected_count += 1 extra_hashes -= 1 - balanced_hashes[resolution] = hashes[:selected_count] - selected_total = sum(len(hashes) for hashes in balanced_hashes.values()) if selected_total < max_results: missing_hashes = max_results - selected_total - - for resolution, hashes in hashes_by_resolution.items(): + for resolution, hash_list in hashes_by_resolution.items(): if missing_hashes <= 0: break - current_count = len(balanced_hashes[resolution]) - available_hashes = hashes[current_count : current_count + missing_hashes] + available_hashes = hash_list[current_count:current_count + missing_hashes] balanced_hashes[resolution].extend(available_hashes) missing_hashes -= len(available_hashes) diff --git a/comet/utils/models.py b/comet/utils/models.py index 5800ee6..fefe62a 100644 --- a/comet/utils/models.py +++ b/comet/utils/models.py @@ -48,10 +48,10 @@ class BestOverallRanking(BaseRankingModel): webdl: int = 90 -rtn_settings: SettingsModel = SettingsModel() -rtn_ranking: BestOverallRanking = BestOverallRanking() +rtn_settings = SettingsModel() +rtn_ranking = BestOverallRanking() # For use anywhere -rtn: RTN = RTN(settings=rtn_settings, ranking_model=rtn_ranking) -settings: AppSettings = AppSettings() +rtn = RTN(settings=rtn_settings, ranking_model=rtn_ranking) +settings = AppSettings() database = Database(f"sqlite:///{settings.DATABASE_PATH}")