Skip to content

Commit

Permalink
fix: 0.30ms faster :pro:
Browse files Browse the repository at this point in the history
  • Loading branch information
@g0ldyy
g0ldyy committed Jul 3, 2024
1 parent 53e8d5e commit 6b1daa6
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 70 deletions.
48 changes: 21 additions & 27 deletions comet/api/stream.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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']}",
}
)

Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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']}",
}
)

Expand Down
62 changes: 23 additions & 39 deletions comet/utils/general.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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)

Expand Down
8 changes: 4 additions & 4 deletions comet/utils/models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}")

0 comments on commit 6b1daa6

Please sign in to comment.