mirror of
https://github.com/g0ldyy/comet.git
synced 2026-01-12 01:16:12 +01:00
Merge pull request #461 from g0ldyy/feat/kitsu-offset
Add Kitsu-to-IMDB anime mapping infrastructure
This commit is contained in:
@@ -12,6 +12,7 @@ from comet.debrid.exceptions import DebridAuthError
|
||||
from comet.debrid.manager import get_debrid_extension
|
||||
from comet.metadata.filter import release_filter
|
||||
from comet.metadata.manager import MetadataScraper
|
||||
from comet.services.anime import anime_mapper
|
||||
from comet.services.debrid import DebridService
|
||||
from comet.services.lock import DistributedLock, is_scrape_in_progress
|
||||
from comet.services.orchestration import TorrentManager
|
||||
@@ -318,6 +319,28 @@ async def stream(
|
||||
get_client_ip(request),
|
||||
)
|
||||
|
||||
is_kitsu = media_id.startswith("kitsu:")
|
||||
search_episode = episode
|
||||
search_season = season
|
||||
|
||||
if is_kitsu and episode is not None:
|
||||
kitsu_mapping = anime_mapper.get_kitsu_episode_mapping(id)
|
||||
if kitsu_mapping:
|
||||
from_episode = kitsu_mapping.get("from_episode")
|
||||
from_season = kitsu_mapping.get("from_season")
|
||||
if from_episode:
|
||||
new_episode = from_episode + episode - 1
|
||||
if new_episode != episode:
|
||||
search_episode = new_episode
|
||||
|
||||
if from_season and from_season != season:
|
||||
search_season = from_season
|
||||
if search_season != season or search_episode != episode:
|
||||
logger.log(
|
||||
"SCRAPER",
|
||||
f"📺 Multi-part anime detected (kitsu:{id}): searching for S{search_season:02d}E{search_episode:02d} instead of S{season:02d}E{episode:02d}",
|
||||
)
|
||||
|
||||
torrent_manager = TorrentManager(
|
||||
debrid_service,
|
||||
config["debridApiKey"],
|
||||
@@ -332,7 +355,9 @@ async def stream(
|
||||
episode,
|
||||
aliases,
|
||||
settings.REMOVE_ADULT_CONTENT and config["removeTrash"],
|
||||
is_kitsu=media_id.startswith("kitsu:"),
|
||||
is_kitsu=is_kitsu,
|
||||
search_episode=search_episode,
|
||||
search_season=search_season,
|
||||
)
|
||||
|
||||
await torrent_manager.get_cached_torrents()
|
||||
|
||||
@@ -370,6 +370,18 @@ async def setup_database():
|
||||
"""
|
||||
)
|
||||
|
||||
await database.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS kitsu_imdb_mapping (
|
||||
kitsu_id TEXT PRIMARY KEY,
|
||||
imdb_id TEXT,
|
||||
title TEXT,
|
||||
from_season INTEGER,
|
||||
from_episode INTEGER
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
await database.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS digital_release_cache (
|
||||
|
||||
+121
-13
@@ -45,8 +45,11 @@ class AnimeMapper:
|
||||
self._refresh_task = None
|
||||
|
||||
self.anime_imdb_ids = set()
|
||||
self._kitsu_mapping_cache = {}
|
||||
|
||||
self._aod_url = "https://github.com/manami-project/anime-offline-database/releases/latest/download/anime-offline-database-minified.json"
|
||||
self._fribb_url = "https://raw.githubusercontent.com/Fribb/anime-lists/refs/heads/master/anime-list-full.json"
|
||||
self._kitsu_imdb_url = "https://raw.githubusercontent.com/TheBeastLT/stremio-kitsu-anime/master/static/data/imdb_mapping.json"
|
||||
|
||||
async def load_anime_mapping(self, session: aiohttp.ClientSession | None = None):
|
||||
if not settings.ANIME_MAPPING_ENABLED:
|
||||
@@ -58,15 +61,24 @@ class AnimeMapper:
|
||||
count = await database.fetch_val("SELECT COUNT(*) FROM anime_entries")
|
||||
if count and count > 0:
|
||||
await self._load_provider_ids()
|
||||
await self._load_kitsu_mapping_cache()
|
||||
|
||||
if await self._is_cache_stale():
|
||||
kitsu_count = await database.fetch_val(
|
||||
"SELECT COUNT(*) FROM kitsu_imdb_mapping"
|
||||
)
|
||||
needs_kitsu_refresh = (
|
||||
kitsu_count == 0 or len(self._kitsu_mapping_cache) == 0
|
||||
)
|
||||
|
||||
if await self._is_cache_stale() or needs_kitsu_refresh:
|
||||
self._refresh_task = asyncio.create_task(
|
||||
self._refresh_from_remote(background=True)
|
||||
)
|
||||
|
||||
self.loaded = True
|
||||
logger.log(
|
||||
"COMET", f"✅ Anime mapping loaded from database: {count} entries"
|
||||
"COMET",
|
||||
f"✅ Anime mapping loaded from database: {count} entries, {len(self._kitsu_mapping_cache)} Kitsu-IMDB mappings",
|
||||
)
|
||||
return True
|
||||
|
||||
@@ -148,6 +160,12 @@ class AnimeMapper:
|
||||
|
||||
return val
|
||||
|
||||
def get_kitsu_episode_mapping(self, kitsu_id: str | int):
|
||||
if not self.loaded:
|
||||
return None
|
||||
|
||||
return self._kitsu_mapping_cache.get(str(kitsu_id))
|
||||
|
||||
def is_loaded(self):
|
||||
return self.loaded
|
||||
|
||||
@@ -187,6 +205,29 @@ class AnimeMapper:
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load anime provider IDs: {e}")
|
||||
|
||||
async def _load_kitsu_mapping_cache(self):
|
||||
try:
|
||||
rows = await database.fetch_all(
|
||||
"""
|
||||
SELECT kitsu_id, imdb_id, from_season, from_episode
|
||||
FROM kitsu_imdb_mapping
|
||||
WHERE (from_episode IS NOT NULL AND from_episode > 1)
|
||||
OR from_season IS NOT NULL
|
||||
"""
|
||||
)
|
||||
|
||||
self._kitsu_mapping_cache.clear()
|
||||
for row in rows:
|
||||
kitsu_id = row["kitsu_id"]
|
||||
|
||||
self._kitsu_mapping_cache[str(kitsu_id)] = {
|
||||
"imdb_id": row["imdb_id"],
|
||||
"from_season": row["from_season"],
|
||||
"from_episode": row["from_episode"],
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load Kitsu-IMDB mapping cache: {e}")
|
||||
|
||||
async def _refresh_from_remote(
|
||||
self,
|
||||
session: aiohttp.ClientSession | None = None,
|
||||
@@ -205,7 +246,7 @@ class AnimeMapper:
|
||||
try:
|
||||
logger.log(
|
||||
"COMET",
|
||||
"Downloading anime mapping (Source 1/2: Anime Offline Database)...",
|
||||
"Downloading anime mapping (Source 1/3: Anime Offline Database)...",
|
||||
)
|
||||
async with session.get(self._aod_url) as response_aod:
|
||||
if response_aod.status != 200:
|
||||
@@ -215,7 +256,7 @@ class AnimeMapper:
|
||||
|
||||
logger.log(
|
||||
"COMET",
|
||||
"Downloading anime mapping (Source 2/2: Fribb Anime List)...",
|
||||
"Downloading anime mapping (Source 2/3: Fribb Anime List)...",
|
||||
)
|
||||
async with session.get(self._fribb_url) as response_fribb:
|
||||
if response_fribb.status != 200:
|
||||
@@ -225,13 +266,26 @@ class AnimeMapper:
|
||||
return False
|
||||
data_fribb = orjson.loads(await response_fribb.read())
|
||||
|
||||
anime_list = data_aod.get("data", [])
|
||||
total_entries, total_fribb = await self._persist_mapping(
|
||||
anime_list, data_fribb
|
||||
logger.log(
|
||||
"COMET",
|
||||
"Downloading anime mapping (Source 3/3: Kitsu-IMDB Mapping)...",
|
||||
)
|
||||
async with session.get(self._kitsu_imdb_url) as response_kitsu:
|
||||
if response_kitsu.status != 200:
|
||||
logger.warning(
|
||||
f"Failed to load Kitsu-IMDB mapping: HTTP {response_kitsu.status}"
|
||||
)
|
||||
return False
|
||||
data_kitsu_imdb = orjson.loads(await response_kitsu.read())
|
||||
|
||||
anime_list = data_aod.get("data", [])
|
||||
total_entries = await self._persist_mapping(anime_list, data_fribb)
|
||||
|
||||
await self._persist_kitsu_imdb_mapping(data_kitsu_imdb)
|
||||
|
||||
del data_aod
|
||||
del data_fribb
|
||||
del data_kitsu_imdb
|
||||
del anime_list
|
||||
gc.collect()
|
||||
|
||||
@@ -249,11 +303,12 @@ class AnimeMapper:
|
||||
pass
|
||||
|
||||
await self._load_provider_ids()
|
||||
await self._load_kitsu_mapping_cache()
|
||||
|
||||
self.loaded = True
|
||||
logger.log(
|
||||
"COMET",
|
||||
f"✅ Anime mapping loaded: {total_entries} entries",
|
||||
f"✅ Anime mapping loaded: {total_entries} entries, {len(self._kitsu_mapping_cache)} Kitsu-IMDB mappings cached",
|
||||
)
|
||||
|
||||
return True
|
||||
@@ -272,7 +327,6 @@ class AnimeMapper:
|
||||
ids_batch = []
|
||||
lookup_map = {}
|
||||
total_entries = 0
|
||||
total_fribb = 0
|
||||
|
||||
entries_query = "INSERT INTO anime_entries (id, data) VALUES (:id, :data)"
|
||||
ids_query = """
|
||||
@@ -364,12 +418,10 @@ class AnimeMapper:
|
||||
|
||||
if len(fribb_batch) >= _DB_CHUNK_SIZE:
|
||||
await database.execute_many(ids_query, fribb_batch)
|
||||
total_fribb += len(fribb_batch)
|
||||
fribb_batch.clear()
|
||||
|
||||
if fribb_batch:
|
||||
await database.execute_many(ids_query, fribb_batch)
|
||||
total_fribb += len(fribb_batch)
|
||||
fribb_batch.clear()
|
||||
|
||||
del fribb_batch
|
||||
@@ -384,10 +436,66 @@ class AnimeMapper:
|
||||
{"timestamp": timestamp},
|
||||
)
|
||||
|
||||
return total_entries, total_fribb
|
||||
return total_entries
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to persist anime mapping cache: {exc}")
|
||||
return 0, 0
|
||||
return 0
|
||||
|
||||
async def _persist_kitsu_imdb_mapping(self, kitsu_imdb_data: dict):
|
||||
total_count = 0
|
||||
batch = []
|
||||
batch_size = 1000
|
||||
|
||||
try:
|
||||
async with database.transaction():
|
||||
await database.execute("DELETE FROM kitsu_imdb_mapping")
|
||||
|
||||
insert_query = """
|
||||
INSERT INTO kitsu_imdb_mapping
|
||||
(kitsu_id, imdb_id, title, from_season, from_episode)
|
||||
VALUES (:kitsu_id, :imdb_id, :title, :from_season, :from_episode)
|
||||
ON CONFLICT (kitsu_id) DO UPDATE SET
|
||||
imdb_id = :imdb_id,
|
||||
title = :title,
|
||||
from_season = :from_season,
|
||||
from_episode = :from_episode
|
||||
"""
|
||||
|
||||
for kitsu_id, entry in kitsu_imdb_data.items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
|
||||
imdb_id = entry.get("imdb_id")
|
||||
if not imdb_id:
|
||||
continue
|
||||
|
||||
from_season = entry.get("fromSeason")
|
||||
from_episode = entry.get("fromEpisode")
|
||||
|
||||
batch.append(
|
||||
{
|
||||
"kitsu_id": str(kitsu_id),
|
||||
"imdb_id": imdb_id,
|
||||
"title": entry.get("title"),
|
||||
"from_season": from_season,
|
||||
"from_episode": from_episode,
|
||||
}
|
||||
)
|
||||
|
||||
if len(batch) >= batch_size:
|
||||
await database.execute_many(insert_query, batch)
|
||||
total_count += len(batch)
|
||||
batch.clear()
|
||||
|
||||
if batch:
|
||||
await database.execute_many(insert_query, batch)
|
||||
total_count += len(batch)
|
||||
batch.clear()
|
||||
|
||||
return total_count
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to persist Kitsu-IMDB mapping: {exc}")
|
||||
return 0
|
||||
|
||||
|
||||
anime_mapper = AnimeMapper()
|
||||
|
||||
@@ -28,8 +28,10 @@ class TorrentManager:
|
||||
episode: int,
|
||||
aliases: dict,
|
||||
remove_adult_content: bool,
|
||||
is_kitsu: bool = False, # Kitsu treats each season as separate anime
|
||||
context: str = "live", # "live" or "background"
|
||||
is_kitsu: bool = False,
|
||||
context: str = "live",
|
||||
search_episode: int | None = None,
|
||||
search_season: int | None = None,
|
||||
):
|
||||
self.debrid_service = debrid_service
|
||||
self.debrid_api_key = debrid_api_key
|
||||
@@ -42,6 +44,8 @@ class TorrentManager:
|
||||
self.year_end = year_end
|
||||
self.season = season
|
||||
self.episode = episode
|
||||
self.search_episode = search_episode if search_episode is not None else episode
|
||||
self.search_season = search_season if search_season is not None else season
|
||||
self.aliases = aliases
|
||||
self.remove_adult_content = remove_adult_content
|
||||
self.is_kitsu = is_kitsu
|
||||
@@ -64,8 +68,8 @@ class TorrentManager:
|
||||
title=self.title,
|
||||
year=self.year,
|
||||
year_end=self.year_end,
|
||||
season=self.season,
|
||||
episode=self.episode,
|
||||
season=self.search_season,
|
||||
episode=self.search_episode,
|
||||
context=self.context,
|
||||
)
|
||||
|
||||
@@ -81,7 +85,7 @@ class TorrentManager:
|
||||
)
|
||||
|
||||
if self.is_kitsu:
|
||||
if episode is not None and episode != self.episode:
|
||||
if episode is not None and episode != self.search_episode:
|
||||
continue
|
||||
else:
|
||||
if (season is not None and season != self.season) or (
|
||||
@@ -111,7 +115,7 @@ class TorrentManager:
|
||||
""",
|
||||
{
|
||||
"media_id": self.media_only_id,
|
||||
"episode": self.episode,
|
||||
"episode": self.search_episode,
|
||||
},
|
||||
)
|
||||
else:
|
||||
@@ -136,7 +140,8 @@ class TorrentManager:
|
||||
parsed_data = ParsedData(**orjson.loads(row["parsed"]))
|
||||
|
||||
if row["episode"] is None and parsed_data.episodes:
|
||||
if self.episode not in parsed_data.episodes:
|
||||
target_episode = self.search_episode if self.is_kitsu else self.episode
|
||||
if target_episode not in parsed_data.episodes:
|
||||
continue
|
||||
|
||||
info_hash = row["info_hash"]
|
||||
|
||||
Reference in New Issue
Block a user