Skip to content

Commit

Permalink
fix: faster (we title match check before doing the rest)
Browse files Browse the repository at this point in the history
  • Loading branch information
g0ldyy committed Jul 18, 2024
1 parent 784a088 commit 19fc235
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 13 deletions.
52 changes: 40 additions & 12 deletions comet/api/stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
get_debrid_extension,
get_indexer_manager,
get_zilean,
filter,
get_torrent_hash,
translate,
get_balanced_hashes,
Expand Down Expand Up @@ -143,7 +144,9 @@ async def stream(request: Request, b64config: str, type: str, id: str):
"torrentTitle": data["torrent_title"]
if "torrent_title" in data
else None,
"torrentSize": data["torrent_size"] if "torrent_size" in data else None,
"torrentSize": data["torrent_size"]
if "torrent_size" in data
else None,
"url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{data['index']}",
}
)
Expand Down Expand Up @@ -206,6 +209,31 @@ async def stream(request: Request, b64config: str, type: str, id: str):
f"{len(torrents)} torrents found for {log_name} with {indexer_manager_type}{' and Zilean' if settings.ZILEAN_URL else ''}"
)

if len(torrents) == 0:
return {"streams": []}

indexed_torrents = [(i, torrents[i]["Title"]) for i in range(len(torrents))]
chunk_size = 50
chunks = [
indexed_torrents[i : i + chunk_size]
for i in range(0, len(indexed_torrents), chunk_size)
]

tasks = []
for chunk in chunks:
tasks.append(filter(chunk, name))

filtered_torrents = await asyncio.gather(*tasks)
index_less = 0
for result in filtered_torrents:
for filtered in result:
if not filtered[1]:
del torrents[filtered[0] - index_less]
index_less += 1
continue

logger.info(f"{len(torrents)} torrents passed title match check for {log_name}")

if len(torrents) == 0:
return {"streams": []}

Expand Down Expand Up @@ -237,12 +265,13 @@ async def stream(request: Request, b64config: str, type: str, id: str):

ranked_files = set()
for hash in files:
try:
ranked_file = rtn.rank(
files[hash]["title"], hash, correct_title=name, remove_trash=True
)
except:
continue
# try:
ranked_file = rtn.rank(
files[hash]["title"],
hash, # , correct_title=name, remove_trash=True
)
# except:
# continue

ranked_files.add(ranked_file)

Expand All @@ -259,10 +288,7 @@ async def stream(request: Request, b64config: str, type: str, id: str):
key: (value.model_dump() if isinstance(value, Torrent) else value)
for key, value in sorted_ranked_files.items()
}
torrents_by_hash = {
torrent["InfoHash"]: torrent
for torrent in torrents
}
torrents_by_hash = {torrent["InfoHash"]: torrent for torrent in torrents}
for hash in sorted_ranked_files: # needed for caching
sorted_ranked_files[hash]["data"]["title"] = files[hash]["title"]
sorted_ranked_files[hash]["data"]["torrent_title"] = torrents_by_hash[hash][
Expand All @@ -272,7 +298,9 @@ async def stream(request: Request, b64config: str, type: str, id: str):
"Tracker"
]
sorted_ranked_files[hash]["data"]["size"] = files[hash]["size"]
sorted_ranked_files[hash]["data"]["torrent_size"] = torrents_by_hash[hash]["Size"]
sorted_ranked_files[hash]["data"]["torrent_size"] = torrents_by_hash[hash][
"Size"
]
sorted_ranked_files[hash]["data"]["index"] = files[hash]["index"]

json_data = json.dumps(sorted_ranked_files).replace("'", "''")
Expand Down
21 changes: 20 additions & 1 deletion comet/utils/general.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
import aiohttp
import bencodepy

from RTN import parse, title_match

from comet.utils.logger import logger
from comet.utils.models import settings, ConfigModel

Expand Down Expand Up @@ -247,7 +249,9 @@ async def get_indexer_manager(
response = await response.json()

for result in response:
result["InfoHash"] = result["infoHash"] if "infoHash" in result else None
result["InfoHash"] = (
result["infoHash"] if "infoHash" in result else None
)
result["Title"] = result["title"]
result["Size"] = result["size"]
result["Link"] = result["downloadUrl"]
Expand Down Expand Up @@ -313,6 +317,21 @@ async def get_zilean(
return results


async def filter(torrents: list, name: str):
results = []
for torrent in torrents:
index = torrent[0]
torrent = torrent[1]

if title_match(name, parse(torrent).parsed_title):
results.append((index, True))
continue

results.append((index, False))

return results


async def get_torrent_hash(session: aiohttp.ClientSession, torrent: tuple):
index = torrent[0]
torrent = torrent[1]
Expand Down

0 comments on commit 19fc235

Please sign in to comment.