Merge pull request #461 from g0ldyy/feat/kitsu-offset

Add Kitsu-to-IMDB anime mapping infrastructure
This commit is contained in:
Goldy
2026-01-09 02:02:05 +01:00
committed by GitHub
4 changed files with 171 additions and 21 deletions
+26 -1
View File
@@ -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()
+12
View File
@@ -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
View File
@@ -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()
+12 -7
View File
@@ -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"]