novel-downloader 1.4.5__py3-none-any.whl → 2.0.0__py3-none-any.whl

novel_downloader/core/downloaders/qidian.py

@@ -3,15 +3,22 @@
 novel_downloader.core.downloaders.qidian
 ----------------------------------------
 
+Downloader implementation for Qidian novels,
+with handling for restricted and encrypted chapters
 """
 
 import asyncio
-import json
 from collections.abc import Awaitable, Callable
-from contextlib import suppress
-from typing import Any, cast
+from pathlib import Path
+from typing import Any
 
 from novel_downloader.core.downloaders.base import BaseDownloader
+from novel_downloader.core.downloaders.registry import register_downloader
+from novel_downloader.core.downloaders.signals import (
+    STOP,
+    Progress,
+    StopToken,
+)
 from novel_downloader.core.interfaces import (
     FetcherProtocol,
     ParserProtocol,
@@ -19,23 +26,30 @@ from novel_downloader.core.interfaces import (
 from novel_downloader.models import (
     BookConfig,
     ChapterDict,
-    CidTask,
     DownloaderConfig,
-    HtmlTask,
 )
-from novel_downloader.utils.chapter_storage import ChapterStorage
-from novel_downloader.utils.file_utils import save_as_json, save_as_txt
-from novel_downloader.utils.time_utils import (
-    async_sleep_with_random_delay,
-    calculate_time_difference,
+from novel_downloader.utils import (
+    ChapterStorage,
+    async_jitter_sleep,
 )
 
 
+@register_downloader(site_keys=["qidian", "qd"])
 class QidianDownloader(BaseDownloader):
     """
-    Specialized downloader for Qidian novels.
+    Specialized downloader for Qidian (起点) novels.
+
+    Processes each chapter in a single worker that
+    handles fetch -> parse -> enqueue storage.
     """
 
+    DEFAULT_SOURCE_ID = 0
+    ENCRYPTED_SOURCE_ID = 1
+    PRIORITIES_MAP = {
+        DEFAULT_SOURCE_ID: 0,
+        ENCRYPTED_SOURCE_ID: 1,
+    }
+
     def __init__(
         self,
         fetcher: FetcherProtocol,
@@ -50,6 +64,7 @@ class QidianDownloader(BaseDownloader):
         book: BookConfig,
         *,
         progress_hook: Callable[[int, int], Awaitable[None]] | None = None,
+        cancel_event: asyncio.Event | None = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -58,306 +73,200 @@ class QidianDownloader(BaseDownloader):
         :param book: BookConfig with at least 'book_id'.
         """
         TAG = "[Downloader]"
+        NUM_WORKERS = 1
+
         book_id = book["book_id"]
         start_id = book.get("start_id")
         end_id = book.get("end_id")
         ignore_set = set(book.get("ignore_ids", []))
 
-        raw_base = self.raw_data_dir / book_id
-        cache_base = self.cache_dir / book_id
-        info_path = raw_base / "book_info.json"
-        chapters_html_dir = cache_base / "html"
-
+        raw_base = self._raw_data_dir / book_id
         raw_base.mkdir(parents=True, exist_ok=True)
-        if self.save_html:
-            chapters_html_dir.mkdir(parents=True, exist_ok=True)
-        normal_cs = ChapterStorage(
-            raw_base=raw_base,
-            namespace="chapters",
-            backend_type=self._config.storage_backend,
-            batch_size=self._config.storage_batch_size,
-        )
-        encrypted_cs = ChapterStorage(
+        html_dir = self._debug_dir / book_id / "html"
+
+        chapter_storage = ChapterStorage(
             raw_base=raw_base,
-            namespace="encrypted_chapters",
-            backend_type=self._config.storage_backend,
-            batch_size=self._config.storage_batch_size,
+            priorities=self.PRIORITIES_MAP,
         )
-
-        # load or fetch book_info
-        book_info: dict[str, Any]
-        re_fetch = True
-        old_data: dict[str, Any] = {}
-
-        if info_path.exists():
-            try:
-                old_data = json.loads(info_path.read_text("utf-8"))
-                days, *_ = calculate_time_difference(
-                    old_data.get("update_time", ""), "UTC+8"
-                )
-                re_fetch = days > 1
-            except Exception:
-                re_fetch = True
-
-        if re_fetch:
-            info_html = await self.fetcher.get_book_info(book_id)
-            if self.save_html:
-                for i, html in enumerate(info_html):
-                    save_as_txt(html, chapters_html_dir / f"info_{i}.html")
-            book_info = self.parser.parse_book_info(info_html)
-
-            if book_info.get("book_name") != "未找到书名":
-                save_as_json(book_info, info_path)
-            else:
-                self.logger.warning("%s 书籍信息未找到, book_id = %s", TAG, book_id)
-                book_info = old_data or {"book_name": "未找到书名"}
-        else:
-            book_info = old_data
-
-        vols = book_info.get("volumes", [])
-        total_chapters = 0
-        for vol in vols:
-            total_chapters += len(vol.get("chapters", []))
-        if total_chapters == 0:
-            self.logger.warning("%s 书籍没有章节可下载: book_id=%s", TAG, book_id)
-            return
-
-        completed_count = 0
-
-        # setup queue
-        cid_queue: asyncio.Queue[CidTask] = asyncio.Queue()
-        html_queue: asyncio.Queue[HtmlTask] = asyncio.Queue()
-        save_queue: asyncio.Queue[ChapterDict] = asyncio.Queue()
-
-        async def fetcher_worker(
-            book_id: str,
-            cid_queue: asyncio.Queue[CidTask],
-            html_queue: asyncio.Queue[HtmlTask],
-            retry_times: int,
-        ) -> None:
-            while True:
-                task = await cid_queue.get()
-                cid = task.cid
-                if not cid:
-                    self.logger.warning("[Fetcher] Skipped empty cid task: %s", task)
-                    cid_queue.task_done()
-                    continue
-
-                if cid in ignore_set:
-                    cid_queue.task_done()
-                    continue
-
+        chapter_storage.connect()
+
+        def cancelled() -> bool:
+            return bool(cancel_event and cancel_event.is_set())
+
+        try:
+            # ---- metadata ---
+            book_info = await self.load_book_info(book_id=book_id, html_dir=html_dir)
+            if not book_info:
+                return
+
+            vols = book_info["volumes"]
+            total_chapters = sum(len(v["chapters"]) for v in vols)
+            if total_chapters == 0:
+                self.logger.warning("%s 书籍没有章节可下载: %s", TAG, book_id)
+                return
+
+            progress = Progress(total_chapters, progress_hook)
+
+            # ---- queues & batching ---
+            cid_q: asyncio.Queue[str | StopToken] = asyncio.Queue()
+            save_q: asyncio.Queue[ChapterDict | StopToken] = asyncio.Queue()
+            default_batch: list[ChapterDict] = []
+            encrypted_batch: list[ChapterDict] = []
+
+            def select_batch(chap: ChapterDict) -> tuple[list[ChapterDict], int]:
+                # set extra.encrypted (by parser); default to plain if absent.
+                if chap.get("extra", {}).get("encrypted", False):
+                    return encrypted_batch, self.ENCRYPTED_SOURCE_ID
+                return default_batch, self.DEFAULT_SOURCE_ID
+
+            async def flush_batch(batch: list[ChapterDict], src: int) -> None:
+                if not batch:
+                    return
                 try:
-                    html_list = await self.fetcher.get_book_chapter(book_id, cid)
-                    await html_queue.put(
-                        HtmlTask(cid=cid, retry=task.retry, html_list=html_list)
-                    )
-                    self.logger.info("[Fetcher] Downloaded chapter %s", cid)
-                    await async_sleep_with_random_delay(
-                        self.request_interval,
-                        mul_spread=1.1,
-                        max_sleep=self.request_interval + 2,
-                    )
-
+                    chapter_storage.upsert_chapters(batch, src)
                 except Exception as e:
-                    if task.retry < retry_times:
-                        await cid_queue.put(
-                            CidTask(
-                                prev_cid=task.prev_cid,
-                                cid=cid,
-                                retry=task.retry + 1,
-                            )
-                        )
-                        self.logger.info(
-                            "[Fetcher] Re-queued chapter %s for retry #%d: %s",
-                            cid,
-                            task.retry + 1,
-                            e,
-                        )
-                        backoff = self.backoff_factor * (2**task.retry)
-                        await async_sleep_with_random_delay(
-                            base=backoff,
-                            mul_spread=1.2,
-                            max_sleep=backoff + 3,
-                        )
-                    else:
-                        self.logger.warning(
-                            "[Fetcher] Max retries reached for chapter %s: %s",
-                            cid,
-                            e,
-                        )
-
-                finally:
-                    cid_queue.task_done()
-
-        async def parser_worker(
-            cid_queue: asyncio.Queue[CidTask],
-            html_queue: asyncio.Queue[HtmlTask],
-            save_queue: asyncio.Queue[ChapterDict],
-            retry_times: int,
-        ) -> None:
-            while True:
-                task = await html_queue.get()
-                skip_retry = False
-                try:
-                    chap_json: ChapterDict | None = None
-                    if self.check_restricted(task.html_list):
-                        self.logger.info(
-                            "[Parser] Skipped restricted page for cid %s", task.cid
-                        )
-                        skip_retry = True
-                        raise ValueError("Restricted content detected")
-
-                    is_encrypted = self.check_encrypted(task.html_list)
-                    chap_json = await asyncio.to_thread(
-                        self.parser.parse_chapter,
-                        task.html_list,
-                        task.cid,
+                    self.logger.error(
+                        "[Storage] batch upsert failed (size=%d, src=%d): %s",
+                        len(batch),
+                        src,
+                        e,
+                        exc_info=True,
                     )
-                    if is_encrypted:
-                        skip_retry = True
-                    if self.save_html:
-                        folder = chapters_html_dir / (
-                            "html_encrypted" if is_encrypted else "html_plain"
-                        )
-                        html_path = folder / f"{task.cid}.html"
-                        save_as_txt(task.html_list[0], html_path, on_exist="skip")
-                        self.logger.debug(
-                            "%s Saved raw HTML for chapter %s to %s",
-                            TAG,
-                            task.cid,
-                            html_path,
-                        )
-                    if chap_json:
-                        await save_queue.put(chap_json)
-                        self.logger.info(
-                            "[Parser] saved chapter %s",
-                            task.cid,
-                        )
-                    else:
-                        raise ValueError("Empty parse result")
-                except Exception as e:
-                    if not skip_retry and task.retry < retry_times:
-                        await cid_queue.put(
-                            CidTask(prev_cid=None, cid=task.cid, retry=task.retry + 1)
-                        )
-                        self.logger.info(
-                            "[Parser] Re-queued cid %s for retry #%d: %s",
-                            task.cid,
-                            task.retry + 1,
-                            e,
-                        )
-                    elif not skip_retry:
-                        self.logger.warning(
-                            "[Parser] Max retries reached for cid %s: %s",
-                            task.cid,
-                            e,
-                        )
+                else:
+                    await progress.bump(len(batch))
                 finally:
-                    html_queue.task_done()
-
-        async def storage_worker(
-            normal_cs: ChapterStorage,
-            encrypted_cs: ChapterStorage,
-            save_queue: asyncio.Queue[ChapterDict],
-        ) -> None:
-            nonlocal completed_count
-            while True:
-                item = await save_queue.get()
-                try:
-                    is_encrypted = item.get("extra", {}).get("encrypted", False)
-                    cs = encrypted_cs if is_encrypted else normal_cs
-                    cs.save(cast(ChapterDict, item))
-                    completed_count += 1
-                    if progress_hook:
-                        await progress_hook(completed_count, total_chapters)
-                except Exception as e:
-                    self.logger.error("[storage_worker] Failed to save: %s", e)
-                finally:
-                    save_queue.task_done()
-
-        fetcher_task = asyncio.create_task(
-            fetcher_worker(
-                book_id,
-                cid_queue,
-                html_queue,
-                self.retry_times,
-            )
-        )
-
-        parser_task = asyncio.create_task(
-            parser_worker(
-                cid_queue,
-                html_queue,
-                save_queue,
-                self.retry_times,
-            )
-        )
-
-        storage_task = asyncio.create_task(
-            storage_worker(
-                normal_cs=normal_cs,
-                encrypted_cs=encrypted_cs,
-                save_queue=save_queue,
-            )
-        )
-
-        found_start = start_id is None
-        stop_early = False
-
-        for vol in book_info.get("volumes", []):
-            chapters = vol.get("chapters", [])
-            for chap in chapters:
-                if stop_early:
-                    break
-
-                cid = chap.get("chapterId")
-                if not cid:
-                    continue
-
-                if not found_start:
-                    if cid == start_id:
-                        found_start = True
-                    else:
-                        completed_count += 1
+                    batch.clear()
+
+            async def flush_all() -> None:
+                await flush_batch(default_batch, self.DEFAULT_SOURCE_ID)
+                await flush_batch(encrypted_batch, self.ENCRYPTED_SOURCE_ID)
+
+            # ---- workers ---
+            sem = asyncio.Semaphore(self.workers)
+
+            async def storage_worker() -> None:
+                """
+                Consumes parsed chapters, batches by source, flushes on threshold.
+
+                Terminates after receiving STOP from each chapter worker.
+
+                On cancel: drains queue, flushes once, then waits for remaining STOPs.
+                """
+                stop_count = 0
+                while True:
+                    chap = await save_q.get()
+                    if isinstance(chap, StopToken):
+                        stop_count += 1
+                        if stop_count == NUM_WORKERS:
+                            await flush_all()
+                            return
                         continue
 
-                if end_id is not None and cid == end_id:
-                    stop_early = True
-
-                if cid in ignore_set:
-                    continue
-
-                if normal_cs.exists(cid) and self.skip_existing:
-                    completed_count += 1
-                    continue
+                    batch, src = select_batch(chap)
+                    batch.append(chap)
+                    if len(batch) >= self.storage_batch_size:
+                        await flush_batch(batch, src)
+
+                    if cancelled():
+                        # Drain whatever is already parsed
+                        try:
+                            while True:
+                                nxt = save_q.get_nowait()
+                                if isinstance(nxt, StopToken):
+                                    stop_count += 1
+                                else:
+                                    nbatch, nsrc = select_batch(nxt)
+                                    nbatch.append(nxt)
+                        except asyncio.QueueEmpty:
+                            pass
+                        await flush_all()
+                        # Wait for remaining STOPs to arrive
+                        while stop_count < NUM_WORKERS:
+                            nxt = await save_q.get()
+                            if nxt is STOP:
+                                stop_count += 1
+                        return
+
+            async def chapter_worker() -> None:
+                """
+                Single worker: fetch + parse with retry, then enqueue ChapterDict.
+
+                Exits on STOP. If cancelled, does not start a new fetch; signals STOP.
+                """
+                while True:
+                    cid = await cid_q.get()
+                    if isinstance(cid, StopToken):
+                        await save_q.put(STOP)
+                        return
+
+                    if not cid or cid in ignore_set:
+                        continue
 
-                await cid_queue.put(CidTask(cid=cid, prev_cid=None))
+                    if cancelled():
+                        await save_q.put(STOP)
+                        return
 
-            if stop_early:
-                break
+                    async with sem:
+                        chap = await self._process_chapter(book_id, cid, html_dir)
+                    if chap and not cancelled():
+                        await save_q.put(chap)
 
-        await cid_queue.join()
-        await html_queue.join()
-        await save_queue.join()
+                    await async_jitter_sleep(
+                        self.request_interval,
+                        mul_spread=1.1,
+                        max_sleep=self.request_interval + 2,
+                    )
 
-        for task in [fetcher_task, parser_task, storage_task]:
-            task.cancel()
-            with suppress(asyncio.CancelledError):
-                await task
+            async def producer() -> None:
+                """
+                Enqueue chapter IDs respecting start/end/skip_existing.
 
-        normal_cs.close()
-        encrypted_cs.close()
+                Always emits STOP x NUM_WORKERS at the end (even if cancelled early).
+                """
+                try:
+                    async for cid in self._chapter_ids(vols, start_id, end_id):
+                        if cancelled():
+                            break
+                        if self.skip_existing and (
+                            chapter_storage.exists(cid, self.DEFAULT_SOURCE_ID)
+                            or chapter_storage.exists(cid, self.ENCRYPTED_SOURCE_ID)
+                        ):
+                            # Already have either variant; count as done.
+                            await progress.bump(1)
+                        else:
+                            await cid_q.put(cid)
+                finally:
+                    for _ in range(NUM_WORKERS):
+                        await cid_q.put(STOP)
+
+            # ---- run tasks ---
+            async with asyncio.TaskGroup() as tg:
+                tg.create_task(storage_worker())
+                for _ in range(NUM_WORKERS):
+                    tg.create_task(chapter_worker())
+                tg.create_task(producer())
+
+            # ---- done ---
+            if cancelled():
+                self.logger.info(
+                    "%s Novel '%s' cancelled: flushed %d/%d chapters.",
+                    TAG,
+                    book_info.get("book_name", "unknown"),
+                    progress.done,
+                    progress.total,
+                )
+            else:
+                self.logger.info(
+                    "%s Novel '%s' download completed.",
+                    TAG,
+                    book_info.get("book_name", "unknown"),
+                )
 
-        self.logger.info(
-            "%s Novel '%s' download completed.",
-            TAG,
-            book_info.get("book_name", "unknown"),
-        )
-        return
+        finally:
+            chapter_storage.close()
 
     @staticmethod
-    def check_restricted(html_list: list[str]) -> bool:
+    def _check_restricted(html_list: list[str]) -> bool:
         """
         Return True if page content indicates access restriction
         (e.g. not subscribed/purchased).
@@ -370,7 +279,58 @@ class QidianDownloader(BaseDownloader):
         return any(m in html_list[0] for m in markers)
 
     @staticmethod
-    def check_encrypted(html_list: list[str]) -> bool:
+    def _check_encrypted(html_list: list[str]) -> bool:
         if not html_list:
             return True
         return '"cES":2' in html_list[0]
+
+    async def _process_chapter(
+        self,
+        book_id: str,
+        cid: str,
+        html_dir: Path,
+    ) -> ChapterDict | None:
+        """
+        Fetch, debug-save, parse a single chapter with retries.
+
+        :return: ChapterDict on success, or None on failure.
+        """
+        for attempt in range(self.retry_times + 1):
+            try:
+                html_list = await self.fetcher.get_book_chapter(book_id, cid)
+                if self._check_restricted(html_list):
+                    self.logger.info(
+                        "[ChapterWorker] Restricted content detected: %s", cid
+                    )
+                    return None
+                encrypted = self._check_encrypted(html_list)
+
+                folder = "html_encrypted" if encrypted else "html_plain"
+                self._save_html_pages(html_dir / folder, cid, html_list)
+
+                chap = await asyncio.to_thread(
+                    self.parser.parse_chapter, html_list, cid
+                )
+                if encrypted and not chap:
+                    self.logger.info(
+                        "[ChapterWorker] Fail for encrypted chapter: %s", cid
+                    )
+                    return None
+                if not chap:
+                    raise ValueError("Empty parse result")
+                return chap
+
+            except Exception as e:
+                if attempt < self.retry_times:
+                    self.logger.info(
+                        "[ChapterWorker] Retry %s (%s): %s", cid, attempt + 1, e
+                    )
+                    backoff = self.backoff_factor * (2**attempt)
+                    await async_jitter_sleep(
+                        base=backoff,
+                        mul_spread=1.2,
+                        max_sleep=backoff + 3,
+                    )
+                else:
+                    self.logger.warning("[ChapterWorker] Failed %s: %s", cid, e)
+        return None

novel_downloader/core/downloaders/registry.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+"""
+novel_downloader.core.downloaders.registry
+------------------------------------------
+
+Registry and factory helpers for creating site-specific or common downloaders
+"""
+
+__all__ = ["register_downloader", "get_downloader"]
+
+from collections.abc import Callable, Sequence
+from typing import TypeVar
+
+from novel_downloader.core.downloaders.common import CommonDownloader
+from novel_downloader.core.interfaces import (
+    DownloaderProtocol,
+    FetcherProtocol,
+    ParserProtocol,
+)
+from novel_downloader.models import DownloaderConfig
+
+DownloaderBuilder = Callable[
+    [FetcherProtocol, ParserProtocol, DownloaderConfig],
+    DownloaderProtocol,
+]
+D = TypeVar("D", bound=DownloaderProtocol)
+_DOWNLOADER_MAP: dict[str, DownloaderBuilder] = {}
+
+
+def register_downloader(
+    site_keys: Sequence[str],
+) -> Callable[[type[D]], type[D]]:
+    """
+    Decorator to register a downloader class under given keys.
+
+    :param site_keys: Sequence of site identifiers
+    :return: A class decorator that populates _DOWNLOADER_MAP.
+    """
+
+    def decorator(cls: type[D]) -> type[D]:
+        for key in site_keys:
+            _DOWNLOADER_MAP[key.lower()] = cls
+        return cls
+
+    return decorator
+
+
+def get_downloader(
+    fetcher: FetcherProtocol,
+    parser: ParserProtocol,
+    site: str,
+    config: DownloaderConfig,
+) -> DownloaderProtocol:
+    """
+    Returns an DownloaderProtocol for the given site.
+
+    :param fetcher: Fetcher implementation
+    :param parser: Parser implementation
+    :param site: Site name (e.g., 'qidian')
+    :param config: Downloader configuration
+
+    :return: An instance of a downloader class
+    """
+    site_key = site.lower()
+    try:
+        downloader_cls = _DOWNLOADER_MAP[site_key]
+    except KeyError:
+        return CommonDownloader(fetcher, parser, config, site_key)
+    return downloader_cls(fetcher, parser, config)