diff --git a/trakt/api.py b/trakt/api.py index 6997484..5aa72e9 100644 --- a/trakt/api.py +++ b/trakt/api.py @@ -20,6 +20,7 @@ PeopleI, RecommendationsI, ScrobbleI, + SearchI, ShowsI, ) @@ -47,6 +48,7 @@ def __init__( people_interface: Optional[Type[PeopleI]] = None, networks_interface: Optional[Type[NetworksI]] = None, comments_interface: Optional[Type[CommentsI]] = None, + search_interface: Optional[Type[SearchI]] = None, recommendations_interface: Optional[Type[RecommendationsI]] = None, scrobble_interface: Optional[Type[ScrobbleI]] = None, user: Optional[TraktCredentials] = None, @@ -83,6 +85,7 @@ def __init__( self.people = (people_interface or PeopleI)(self, Executor) self.networks = (networks_interface or NetworksI)(self, Executor) self.comments = (comments_interface or CommentsI)(self, Executor) + self.search = (search_interface or SearchI)(self, Executor) self.recommendations = (recommendations_interface or RecommendationsI)( self, Executor ) @@ -121,4 +124,5 @@ def _get_executor_paths(self) -> List[SuiteInterface]: self.networks, self.recommendations, self.scrobble, + self.search, ] diff --git a/trakt/core/paths/__init__.py b/trakt/core/paths/__init__.py index 418eea9..15015f1 100644 --- a/trakt/core/paths/__init__.py +++ b/trakt/core/paths/__init__.py @@ -16,3 +16,4 @@ from trakt.core.paths.endpoint_mappings.people import PeopleI from trakt.core.paths.endpoint_mappings.recommendations import RecommendationsI from trakt.core.paths.endpoint_mappings.scrobble import ScrobbleI +from trakt.core.paths.endpoint_mappings.search import SearchI diff --git a/trakt/core/paths/endpoint_mappings/search.py b/trakt/core/paths/endpoint_mappings/search.py new file mode 100644 index 0000000..b680dde --- /dev/null +++ b/trakt/core/paths/endpoint_mappings/search.py @@ -0,0 +1,89 @@ +from typing import Iterable, List, Optional, Union + +from trakt.core.paths.path import Path +from trakt.core.paths.response_structs import SearchResult +from trakt.core.paths.suite_interface import SuiteInterface +from trakt.core.paths.validators import ALL_FILTERS, PerArgValidator + +MEDIA_TYPES = ["movie", "show", "episode", "person", "list"] +ID_TYPES = ["trakt", "imdb", "tmdb", "tvdb"] +POSSIBLE_FIELDS = { + "title", + "tagline", + "overview", + "people", + "translations", + "aliases", + "name", + "biography", + "description", +} + + +class SearchI(SuiteInterface): + name = "search" + + paths = { + "text_query": Path( # type: ignore + "search/!type", + [SearchResult], + extended=["full"], + filters=ALL_FILTERS, + pagination=True, + validators=[ + PerArgValidator("type", lambda t: t in MEDIA_TYPES), + PerArgValidator("query", lambda q: isinstance(q, str) and q), + PerArgValidator( + "fields", lambda f: [x in POSSIBLE_FIELDS for x in f.split(",")] + ), + ], + qargs=["fields"], + ), + "id_lookup": Path( # type: ignore + "search/!id_type/!id", + [SearchResult], + extended=["full"], + filters=ALL_FILTERS, + pagination=True, + validators=[ + PerArgValidator("id", lambda t: isinstance(t, (int, str))), + PerArgValidator("id_type", lambda it: it in ID_TYPES), + PerArgValidator( + "type", lambda f: [x in MEDIA_TYPES for x in f.split(",")] + ), + ], + qargs=["type"], + ), + } + + def text_query( + self, + type: Union[str, List[str]], + query: str, + fields: Optional[Union[str, List[str]]] = None, + **kwargs + ) -> Iterable[SearchResult]: + type = [type] if isinstance(type, str) else type + type = ",".join(type) + req = {"type": type, "query": query} + + if fields: + fields = [fields] if isinstance(fields, str) else fields + req["fields"] = ",".join(fields) + + return self.run("text_query", **kwargs, **req) + + def id_lookup( + self, + id_type: str, + id: Union[str, int], + type: Optional[Union[str, List[str]]] = None, + **kwargs + ) -> Iterable[SearchResult]: + req = {"id_type": type, "id": id} + + if type: + type = [type] if isinstance(type, str) else type + req["type"] = ",".join(type) + + return self.run("id_lookup", **kwargs, **req) diff --git a/trakt/core/paths/response_structs.py b/trakt/core/paths/response_structs.py index 48d5bbb..c68fe94 100644 --- a/trakt/core/paths/response_structs.py +++ b/trakt/core/paths/response_structs.py @@ -283,6 +283,17 @@ class CommentAndItem: season: Optional[Season] = None +@dataclass +class SearchResult: + type: str + score: Optional[float] + movie: Optional[Movie] = None + list: Optional[TraktList] = None + person: Optional[Person] = None + episode: Optional[Episode] = None + show: Optional[Show] = None + + @dataclass class MovieScrobble: id: int