music-kraken-core/music_kraken/download/results.py

122 lines
3.6 KiB
Python
Raw Permalink Normal View History

2023-06-12 12:56:14 +00:00
from typing import Tuple, Type, Dict, List, Generator, Union
from dataclasses import dataclass
2023-05-26 08:11:36 +00:00
from ..objects import DatabaseObject
from ..pages import Page, EncyclopaediaMetallum, Musify
@dataclass
class Option:
index: int
music_object: DatabaseObject
class Results:
2024-05-07 11:34:18 +00:00
def __init__(self, max_items_per_page: int = 10, **kwargs) -> None:
2023-06-12 12:56:14 +00:00
self._by_index: Dict[int, DatabaseObject] = dict()
self._page_by_index: Dict[int: Type[Page]] = dict()
2024-05-07 11:34:18 +00:00
self.max_items_per_page = max_items_per_page
2023-06-12 12:56:14 +00:00
def __iter__(self) -> Generator[DatabaseObject, None, None]:
2024-05-07 11:34:18 +00:00
for option in self.formatted_generator():
if isinstance(option, Option):
yield option.music_object
2024-05-07 11:34:18 +00:00
def formatted_generator(self) -> Generator[Union[Type[Page], Option], None, None]:
self._by_index = dict()
2023-06-12 12:56:14 +00:00
self._page_by_index = dict()
2024-05-07 11:34:18 +00:00
2024-05-08 10:23:16 +00:00
def __len__(self) -> int:
return max(self._by_index.keys())
2024-05-07 11:34:18 +00:00
def __getitem__(self, index: int):
return self._by_index[index]
2023-12-29 15:15:54 +00:00
class SearchResults(Results):
2023-05-26 08:11:36 +00:00
def __init__(
self,
2024-05-07 11:34:18 +00:00
pages: Tuple[Type[Page], ...] = None,
**kwargs,
2023-05-26 08:11:36 +00:00
) -> None:
2024-05-07 11:34:18 +00:00
super().__init__(**kwargs)
2023-06-12 12:56:14 +00:00
self.pages = pages or []
2023-05-26 08:11:36 +00:00
# this would initialize a list for every page, which I don't think I want
# self.results = Dict[Type[Page], List[DatabaseObject]] = {page: [] for page in self.pages}
2023-06-12 12:56:14 +00:00
self.results: Dict[Type[Page], List[DatabaseObject]] = {}
2023-05-26 08:11:36 +00:00
def add(self, page: Type[Page], search_result: List[DatabaseObject]):
"""
adds a list of found music objects to the according page
WARNING: if a page already has search results, they are just gonna be overwritten
"""
2023-05-26 08:11:36 +00:00
self.results[page] = search_result
def get_page_results(self, page: Type[Page]) -> "PageResults":
return PageResults(page, self.results.get(page, []))
2024-05-07 11:34:18 +00:00
def __len__(self) -> int:
return sum(min(self.max_items_per_page, len(results)) for results in self.results.values())
2024-05-07 11:34:18 +00:00
def formatted_generator(self):
super().formatted_generator()
i = 0
for page in self.results:
yield page
j = 0
for option in self.results[page]:
yield Option(i, option)
self._by_index[i] = option
2023-06-12 12:56:14 +00:00
self._page_by_index[i] = page
i += 1
j += 1
2024-05-07 11:34:18 +00:00
if j >= self.max_items_per_page:
break
2024-05-07 11:34:18 +00:00
class GoToResults(Results):
def __init__(self, results: List[DatabaseObject], **kwargs):
self.results: List[DatabaseObject] = results
super().__init__(**kwargs)
def __getitem__(self, index: int):
return self.results[index]
def __len__(self) -> int:
return len(self.results)
def formatted_generator(self):
yield from (Option(i, o) for i, o in enumerate(self.results))
class PageResults(Results):
2024-05-07 11:34:18 +00:00
def __init__(self, page: Type[Page], results: List[DatabaseObject], **kwargs) -> None:
super().__init__(**kwargs)
self.page: Type[Page] = page
self.results: List[DatabaseObject] = results
2024-05-07 11:34:18 +00:00
2024-05-07 11:34:18 +00:00
def formatted_generator(self, max_items_per_page: int = 10):
super().formatted_generator()
i = 0
yield self.page
2023-05-26 08:11:36 +00:00
for option in self.results:
yield Option(i, option)
self._by_index[i] = option
2023-06-12 12:56:14 +00:00
self._page_by_index[i] = self.page
i += 1
2024-05-07 11:34:18 +00:00
def __len__(self) -> int:
return len(self.results)