import asyncio import tempfile import xml.etree.ElementTree as XML from pathlib import Path import aiofiles import aiohttp import questionary from aiohttp import ClientError from rich import print, status from rich.progress import MofNCompleteColumn, Progress, SpinnerColumn, TextColumn from packager import zip_files from sites import Provider class Manga: def __init__(self, provider: Provider, cached: list[float] = []) -> None: self.site = provider self.cached = cached spinner = status.Status("Fetching manga info") spinner.start() self.info, self.chapters = self.site.get_mediainfo() spinner.stop() self.slug = self.slugify(self.info.get("series")) # type: ignore @staticmethod def slugify(string: str) -> str: string = "".join(x for x in string if not x in '<>:"/\\|?*') return str(string) def _generateComicInfo(self, chapter: dict, dir: Path): filename = dir.joinpath("ComicInfo.xml") ComicInfo = XML.Element("ComicInfo") XML.SubElement(ComicInfo, "Series").text = self.info.get("series") XML.SubElement(ComicInfo, "Writer").text = ",".join(self.info.get("writer", [])) # type: ignore XML.SubElement(ComicInfo, "Penciller").text = ",".join(self.info.get("penciller", [])) # type: ignore XML.SubElement(ComicInfo, "Genre").text = ",".join(self.info.get("genre", [])) # type: ignore XML.SubElement(ComicInfo, "Summary").text = self.info.get("summary") XML.SubElement(ComicInfo, "Number").text = str(chapter.get("nr")) XML.SubElement(ComicInfo, "Title").text = chapter.get("title") XML.SubElement(ComicInfo, "LanguageISO").text = self.info.get("languageISO") XML.SubElement(ComicInfo, "PageCount").text = str(len(chapter.get("images"))) # type: ignore XML.SubElement(ComicInfo, "ScanInformation").text = self.info.get( "scanInformation" ) tree = XML.ElementTree(ComicInfo) XML.indent(tree, space="\t", level=0) tree.write(str(filename), encoding="utf-8") return filename @staticmethod def printMangaInfo(info: dict): print(info) def choose_chapters(self, auto): choices = [ questionary.Choice( f"{chapter.get('nr'):g}: {chapter.get('title')}", chapter, checked=chapter.get("nr") not in self.cached, ) for chapter in self.chapters ] self.chapters = ( questionary.checkbox("Choose which chapters to download", choices) .skip_if( auto, [ chapter for chapter in self.chapters if chapter.get("nr") not in self.cached ], ) .ask() ) def download(self, save_dir: Path): if not self.chapters: print("No new [orange3]Chapter[/] were found...") return async def process_all_chapters(): chapter_sem = asyncio.Semaphore(8) image_sem = asyncio.Semaphore(40) connector = aiohttp.TCPConnector(limit=100, ttl_dns_cache=300) async with aiohttp.ClientSession( headers=self.site.headers, connector=connector ) as session: with Progress( TextColumn("[progress.description]{task.description}"), MofNCompleteColumn(), SpinnerColumn(), ) as p: task = p.add_task( f"[bold]Processing [orange3]{self.info.get('series')}[/]", total=len(self.chapters), ) async def download_and_zip(chapter): async with chapter_sem: chapter_nr = chapter.get("nr") chapter_temp = Path(temp_dir) / f"ch_{chapter_nr:g}" chapter_temp.mkdir(parents=True, exist_ok=True) info_file = self._generateComicInfo(chapter, chapter_temp) img_urls = chapter.get("images", []) img_files = await self._async_download_images_internal( img_urls, chapter_temp, session, image_sem ) all_files = [info_file] + [f for f in img_files if f] zip_files(self.slug, chapter, all_files, save_dir) self.cached.append(chapter_nr) p.update(task, advance=1) await asyncio.gather( *(download_and_zip(ch) for ch in self.chapters) ) with tempfile.TemporaryDirectory(prefix="manga_") as temp_dir: asyncio.run(process_all_chapters()) async def _async_download_images_internal(self, urls, temp_dir, session, semaphore): """Worker to manage the image download tasks for a single chapter.""" async def bounded_download(url): async with semaphore: return await self.download_image( session, url, temp_dir, self.site.headers ) tasks = [bounded_download(url) for url in urls] return await asyncio.gather(*tasks) @staticmethod async def download_image( session: aiohttp.ClientSession, url: str, temp_dir: Path, headers: dict, retries=3, timeout=10, ): filename = temp_dir.joinpath(Path(url).name) for attempt in range(1, retries + 1): try: async with session.get(url, headers=headers, timeout=timeout) as response: # type: ignore if response.status == 200: async with aiofiles.open(filename, "wb") as f: await f.write(await response.read()) return filename else: raise ClientError(f"Bad status {response.status} for {url}") except ( asyncio.TimeoutError, ClientError, aiohttp.ClientConnectorError, ) as e: if attempt < retries: await asyncio.sleep(2**attempt) else: print(f"Failed to download {url}: {e}") return None