Write identify function

This commit is contained in:
Zoe Roux 2025-05-05 12:37:29 +02:00
parent a6f48d1f9a
commit 82a5300ba8
No known key found for this signature in database
4 changed files with 65 additions and 10 deletions

View File

@ -9,6 +9,7 @@ if __name__ == "__main__":
from guessit.api import default_api from guessit.api import default_api
from typing import cast, List, Any from typing import cast, List, Any
from rebulk import Rebulk from rebulk import Rebulk
from rebulk.match import MatchesDict
try: try:
from . import rules from . import rules
@ -20,35 +21,41 @@ rblk = cast(Rebulk, default_api.rebulk)
rblk.rules(rules) rblk.rules(rules)
def guessit(name: str, *, xem_titles: List[str] = [], extra_flags: dict[str, Any] = {}): def guessit(
name: str,
*,
expected_titles: List[str] = [],
extra_flags: dict[str, Any] = {},
) -> MatchesDict:
return default_api.guessit( return default_api.guessit(
name, name,
{ {
"episode_prefer_number": True, "episode_prefer_number": True,
"excludes": "language", "excludes": "language",
"expected_title": xem_titles, "expected_title": expected_titles,
"enforce_list": True
} }
| extra_flags, | extra_flags,
) )
# Only used to test localy # Only used to test locally
if __name__ == "__main__": if __name__ == "__main__":
import sys import sys
import json import json
from providers.implementations.thexem import TheXemClient # from providers.implementations.thexem import TheXemClient
from guessit.jsonutils import GuessitEncoder from guessit.jsonutils import GuessitEncoder
from aiohttp import ClientSession from aiohttp import ClientSession
import asyncio import asyncio
async def main(): async def main():
async with ClientSession() as client: async with ClientSession() as client:
xem = TheXemClient(client) # xem = TheXemClient(client)
advanced = any(x == "-a" for x in sys.argv) advanced = any(x == "-a" for x in sys.argv)
ret = guessit( ret = guessit(
sys.argv[1], sys.argv[1],
xem_titles=await xem.get_expected_titles(), expected_titles=[],
extra_flags={"advanced": advanced}, extra_flags={"advanced": advanced},
) )
print(json.dumps(ret, cls=GuessitEncoder, indent=4)) print(json.dumps(ret, cls=GuessitEncoder, indent=4))

View File

@ -0,0 +1,47 @@
from .models.videos import Video, Guess
from .guess.guess import guessit
from typing import Literal
from itertools import zip_longest
async def identify(path: str) -> Video:
raw = guessit(path, expected_titles=[])
# guessit should only return one (according to the doc)
title: str = raw.get("title", [])[0]
kind: Literal["movie"] | Literal["episode"] = raw.get("type", [])[0]
version: int = raw.get("version", [])[0]
# apparently guessit can return multiples but tbh idk what to do with
# multiples part. we'll just ignore them for now
part: int = raw.get("part", [])[0]
years: list[int] = raw.get("year", [])
seasons: list[int] = raw.get("season", [])
episodes: list[int] = raw.get("episode", [])
guess = Guess(
title=title,
kind=kind,
extraKind=None,
years=years,
episodes=[
Guess.Episode(season=s, episode=e)
for s, e in zip_longest(
seasons,
episodes,
fillvalue=seasons[-1] if len(seasons) < len(episodes) else episodes[-1],
)
],
# TODO: add external ids parsing in guessit
external_id={},
from_="guessit",
raw=raw,
)
return Video(
path=path,
rendering="",
part=part,
version=version,
guess=guess,
)

View File

@ -3,7 +3,7 @@ from __future__ import annotations
from ..utils import Model from ..utils import Model
from .extra import ExtraKind from .extra import ExtraKind
from .metadataid import MetadataId, EpisodeId from .metadataid import MetadataId, EpisodeId
from typing import Optional, Literal from typing import Optional, Literal, Any
class Resource(Model): class Resource(Model):
@ -17,16 +17,17 @@ class VideoInfo(Model):
guesses: dict[str, dict[str, Resource]] guesses: dict[str, dict[str, Resource]]
class Guess(Model): class Guess(Model, extra="allow"):
title: str title: str
kind: Literal["episode"] | Literal["movie"] | Literal["extra"] kind: Literal["episode"] | Literal["movie"] | Literal["extra"]
extraKind: Optional[ExtraKind] extraKind: Optional[ExtraKind]
years: list[int] years: list[int]
episodes: list[Guess.Episode] episodes: list[Guess.Episode]
external_id: dict[str, MetadataId | EpisodeId] external_id: dict[str, MetadataId | EpisodeId]
raw: dict[str, Any]
from_: str from_: str
history: list[Guess] history: list[Guess] = []
class Episode(Model): class Episode(Model):
season: Optional[int] season: Optional[int]
@ -68,4 +69,4 @@ class Video(Model):
guess: Guess guess: Guess
for_: Optional[ for_: Optional[
For.Slug | For.ExternalId | For.Movie | For.Episode | For.Order | For.Special For.Slug | For.ExternalId | For.Movie | For.Episode | For.Order | For.Special
] ] = None