novel-downloader 2.0.1__py3-none-any.whl → 2.0.2__py3-none-any.whl

novel_downloader/core/downloaders/common.py

@@ -12,19 +12,9 @@ from pathlib import Path
 from typing import Any
 
 from novel_downloader.core.downloaders.base import BaseDownloader
-from novel_downloader.core.downloaders.signals import (
-    STOP,
-    Progress,
-    StopToken,
-)
-from novel_downloader.models import (
-    BookConfig,
-    ChapterDict,
-)
-from novel_downloader.utils import (
-    ChapterStorage,
-    async_jitter_sleep,
-)
+from novel_downloader.core.downloaders.signals import STOP, Progress, StopToken
+from novel_downloader.models import BookConfig, ChapterDict
+from novel_downloader.utils import ChapterStorage, async_jitter_sleep
 
 
 class CommonDownloader(BaseDownloader):
@@ -41,7 +31,7 @@ class CommonDownloader(BaseDownloader):
         **kwargs: Any,
     ) -> None:
         """
-        Sentinel-based pipeline with graceful cancellation:
+        Sentinel-based pipeline with cancellation:
 
         Producer -> ChapterWorkers -> StorageWorker.
 
@@ -59,181 +49,167 @@ class CommonDownloader(BaseDownloader):
         raw_base.mkdir(parents=True, exist_ok=True)
         html_dir = self._debug_dir / book_id / "html"
 
-        chapter_storage = ChapterStorage(
-            raw_base=raw_base,
-            priorities=self.PRIORITIES_MAP,
-        )
-        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
+        # --- metadata ---
+        book_info = await self.load_book_info(book_id=book_id, html_dir=html_dir)
+        if not book_info:
+            return
 
-            progress = Progress(total_chapters, progress_hook)
+        vols = book_info["volumes"]
+        plan = self._planned_chapter_ids(vols, start_id, end_id, ignore_set)
+        if not plan:
+            self.logger.info("%s nothing to do after filtering: %s", TAG, book_id)
+            return
 
-            # --- queues & batching ---
-            cid_q: asyncio.Queue[str | StopToken] = asyncio.Queue()
-            save_q: asyncio.Queue[ChapterDict | StopToken] = asyncio.Queue()
-            batch: list[ChapterDict] = []
+        progress = Progress(total=len(plan), hook=progress_hook)
 
-            async def flush_batch() -> None:
-                if not batch:
-                    return
-                try:
-                    chapter_storage.upsert_chapters(batch, self.DEFAULT_SOURCE_ID)
-                except Exception as e:
-                    self.logger.error(
-                        "[Storage] batch upsert failed (size=%d): %s",
-                        len(batch),
-                        e,
-                        exc_info=True,
-                    )
-                else:
-                    await progress.bump(len(batch))
-                finally:
-                    batch.clear()
-
-            # --- stage: storage worker ---
-            async def storage_worker() -> None:
-                """
-                Consumes parsed chapters, writes in batches.
-
-                Terminates after receiving STOP from each chapter worker.
-
-                On cancel: keeps consuming (to avoid blocking producers),
-                flushes, and exits once all STOPs are seen.
-                """
-                stop_count = 0
-                while True:
-                    item = await save_q.get()
-                    if isinstance(item, StopToken):
-                        stop_count += 1
-                        if stop_count == self.workers:
-                            # All chapter workers have exited.
-                            await flush_batch()
-                            return
-                        # else keep waiting for remaining STOPs
-                        continue
-
-                    # Normal chapter
-                    batch.append(item)
-                    if len(batch) >= self.storage_batch_size:
-                        await flush_batch()
+        # --- queues & batching ---
+        cid_q: asyncio.Queue[str | StopToken] = asyncio.Queue(maxsize=self._workers * 2)
+        save_q: asyncio.Queue[ChapterDict | StopToken] = asyncio.Queue(
+            maxsize=self._workers * 2
+        )
+        batch: list[ChapterDict] = []
 
-                    if cancelled():
-                        # Drain whatever is already in the queue
-                        try:
-                            while True:
-                                nxt = save_q.get_nowait()
-                                if isinstance(nxt, StopToken):
-                                    stop_count += 1
-                                else:
-                                    batch.append(nxt)
-                        except asyncio.QueueEmpty:
-                            pass
-                        # Final flush of everything
+        async def flush_batch() -> None:
+            if not batch:
+                return
+            try:
+                storage.upsert_chapters(batch, self.DEFAULT_SOURCE_ID)
+            except Exception as e:
+                self.logger.error(
+                    "[Storage] batch upsert failed (size=%d): %s",
+                    len(batch),
+                    e,
+                    exc_info=True,
+                )
+            else:
+                await progress.bump(len(batch))
+            finally:
+                batch.clear()
+
+        # --- stage: storage worker ---
+        async def storage_worker() -> None:
+            """
+            Consumes parsed chapters, writes in batches.
+
+            Terminates after receiving STOP from each chapter worker.
+
+            On cancel: keeps consuming (to avoid blocking producers),
+            flushes, and exits once all STOPs are seen.
+            """
+            stop_count = 0
+            while True:
+                item = await save_q.get()
+                if isinstance(item, StopToken):
+                    stop_count += 1
+                    if stop_count == self._workers:
+                        # All chapter workers have exited.
                         await flush_batch()
-                        # Wait for remaining STOPs so chapter workers can finish.
-                        while stop_count < self.workers:
-                            nxt = await save_q.get()
+                        return
+                    # else keep waiting for remaining STOPs
+                    continue
+
+                # Normal chapter
+                batch.append(item)
+                if len(batch) >= self._storage_batch_size:
+                    await flush_batch()
+
+                if cancelled():
+                    # Drain whatever is already in the queue
+                    try:
+                        while True:
+                            nxt = save_q.get_nowait()
                             if isinstance(nxt, StopToken):
                                 stop_count += 1
-                        return
-
-            # --- stage: chapter worker ---
-            sem = asyncio.Semaphore(self.workers)
+                            else:
+                                batch.append(nxt)
+                    except asyncio.QueueEmpty:
+                        pass
+                    # Final flush of everything
+                    await flush_batch()
+                    # Wait for remaining STOPs so chapter workers can finish.
+                    while stop_count < self._workers:
+                        nxt = await save_q.get()
+                        if isinstance(nxt, StopToken):
+                            stop_count += 1
+                    return
 
-            async def chapter_worker() -> None:
-                """
-                Fetch + parse with retry, then enqueue to save_q.
+        # --- stage: chapter worker ---
+        async def chapter_worker() -> None:
+            """
+            Fetch + parse with retry, then enqueue to save_q.
+
+            Exits on STOP, or early if cancel is set before starting a new fetch.
+            """
+            while True:
+                cid = await cid_q.get()
+                if isinstance(cid, StopToken):
+                    # Propagate one STOP to storage and exit.
+                    await save_q.put(STOP)
+                    return
 
-                Exits on STOP, or early if cancel is set before starting a new fetch.
-                """
-                while True:
-                    cid = await cid_q.get()
-                    if isinstance(cid, StopToken):
-                        # Propagate one STOP to storage and exit.
-                        await save_q.put(STOP)
-                        return
+                # If cancelled, don't start a new network call; let storage finish.
+                if cancelled():
+                    await save_q.put(STOP)
+                    return
 
-                    if not cid or cid in ignore_set:
-                        # Ignore silently and continue.
-                        continue
+                chap = await self._process_chapter(book_id, cid, html_dir)
+                if chap:
+                    await save_q.put(chap)
 
-                    # If cancelled, don't start a new network call; let storage finish.
-                    if cancelled():
-                        await save_q.put(STOP)
-                        return
+                # polite pacing
+                await async_jitter_sleep(
+                    self._request_interval,
+                    mul_spread=1.1,
+                    max_sleep=self._request_interval + 2,
+                )
 
-                    async with sem:
-                        chap = await self._process_chapter(book_id, cid, html_dir)
-                    if chap:
-                        await save_q.put(chap)
+        # --- stage: producer ---
+        async def producer() -> None:
+            """
+            Enqueue chapter IDs (respecting start/end/skip_existing).
 
-                    # polite pacing
-                    await async_jitter_sleep(
-                        self.request_interval,
-                        mul_spread=1.1,
-                        max_sleep=self.request_interval + 2,
-                    )
-
-            # --- stage: producer ---
-            async def producer() -> None:
-                """
-                Enqueue chapter IDs (respecting start/end/skip_existing).
-
-                Always sends STOP x workers at the end (even if cancelled early),
-                so chapter workers can exit deterministically.
-                """
-                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):
-                            # Count as completed but don't enqueue.
-                            await progress.bump(1)
-                        else:
-                            await cid_q.put(cid)
-                finally:
-                    for _ in range(self.workers):
-                        await cid_q.put(STOP)
-
-            # --- run the pipeline ---
+            Always sends STOP x workers at the end (even if cancelled early),
+            so chapter workers can exit deterministically.
+            """
+            try:
+                for cid in plan:
+                    if cancelled():
+                        break
+                    if self._skip_existing and storage.exists(cid):
+                        # Count as completed but don't enqueue.
+                        await progress.bump(1)
+                    else:
+                        await cid_q.put(cid)
+            finally:
+                for _ in range(self._workers):
+                    await cid_q.put(STOP)
+
+        # --- run the pipeline ---
+        with ChapterStorage(raw_base, priorities=self.PRIORITIES_MAP) as storage:
             async with asyncio.TaskGroup() as tg:
                 tg.create_task(storage_worker())
-                for _ in range(self.workers):
+                for _ in range(self._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"),
-                )
-
-        finally:
-            chapter_storage.close()
+        # --- 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"),
+            )
 
     async def _process_chapter(
         self,
@@ -247,7 +223,7 @@ class CommonDownloader(BaseDownloader):
 
         :return: ChapterDict on success, or None on failure.
         """
-        for attempt in range(self.retry_times + 1):
+        for attempt in range(self._retry_times + 1):
             try:
                 html_list = await self.fetcher.get_book_chapter(book_id, cid)
                 self._save_html_pages(html_dir, cid, html_list)
@@ -258,11 +234,11 @@ class CommonDownloader(BaseDownloader):
                     raise ValueError("Empty parse result")
                 return chap
             except Exception as e:
-                if attempt < self.retry_times:
+                if attempt < self._retry_times:
                     self.logger.info(
                         "[ChapterWorker] Retry %s (%s): %s", cid, attempt + 1, e
                     )
-                    backoff = self.backoff_factor * (2**attempt)
+                    backoff = self._backoff_factor * (2**attempt)
                     await async_jitter_sleep(
                         base=backoff, mul_spread=1.2, max_sleep=backoff + 3
                     )

novel_downloader/core/downloaders/qianbi.py

@@ -13,25 +13,13 @@ 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,
-)
+from novel_downloader.core.downloaders.signals import STOP, Progress, StopToken
 from novel_downloader.models import (
     BookConfig,
     BookInfoDict,
     ChapterDict,
-    DownloaderConfig,
-)
-from novel_downloader.utils import (
-    ChapterStorage,
-    async_jitter_sleep,
 )
+from novel_downloader.utils import ChapterStorage, async_jitter_sleep
 
 
 @register_downloader(site_keys=["qianbi"])
@@ -43,16 +31,6 @@ class QianbiDownloader(BaseDownloader):
     each chapter as a unit (fetch -> parse -> enqueue storage).
     """
 
-    DEFAULT_SOURCE_ID = 0
-
-    def __init__(
-        self,
-        fetcher: FetcherProtocol,
-        parser: ParserProtocol,
-        config: DownloaderConfig,
-    ):
-        super().__init__(fetcher, parser, config, "qianbi")
-
     async def _download_one(
         self,
         book: BookConfig,
@@ -77,16 +55,10 @@ class QianbiDownloader(BaseDownloader):
         raw_base.mkdir(parents=True, exist_ok=True)
         html_dir = self._debug_dir / book_id / "html"
 
-        chapter_storage = ChapterStorage(
-            raw_base=raw_base,
-            priorities=self.PRIORITIES_MAP,
-        )
-        chapter_storage.connect()
-
         def cancelled() -> bool:
             return bool(cancel_event and cancel_event.is_set())
 
-        try:
+        with ChapterStorage(raw_base, priorities=self.PRIORITIES_MAP) as storage:
             # --- metadata ---
             book_info = await self.load_book_info(book_id=book_id, html_dir=html_dir)
             if not book_info:
@@ -95,28 +67,32 @@ class QianbiDownloader(BaseDownloader):
             book_info = await self._repair_chapter_ids(
                 book_id,
                 book_info,
-                chapter_storage,
+                storage,
                 html_dir,
             )
 
             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)
+            plan = self._planned_chapter_ids(vols, start_id, end_id, ignore_set)
+            if not plan:
+                self.logger.info("%s nothing to do after filtering: %s", TAG, book_id)
                 return
 
-            progress = Progress(total_chapters, progress_hook)
+            progress = Progress(total=len(plan), hook=progress_hook)
 
             # --- queues & batching ---
-            cid_q: asyncio.Queue[str | StopToken] = asyncio.Queue()
-            save_q: asyncio.Queue[ChapterDict | StopToken] = asyncio.Queue()
+            cid_q: asyncio.Queue[str | StopToken] = asyncio.Queue(
+                maxsize=self._workers * 2
+            )
+            save_q: asyncio.Queue[ChapterDict | StopToken] = asyncio.Queue(
+                maxsize=self._workers * 2
+            )
             batch: list[ChapterDict] = []
 
             async def flush_batch() -> None:
                 if not batch:
                     return
                 try:
-                    chapter_storage.upsert_chapters(batch, self.DEFAULT_SOURCE_ID)
+                    storage.upsert_chapters(batch, self.DEFAULT_SOURCE_ID)
                 except Exception as e:
                     self.logger.error(
                         "[Storage] batch upsert failed (size=%d): %s",
@@ -144,7 +120,7 @@ class QianbiDownloader(BaseDownloader):
                     item = await save_q.get()
                     if isinstance(item, StopToken):
                         stop_count += 1
-                        if stop_count == self.workers:
+                        if stop_count == self._workers:
                             # All chapter workers have exited.
                             await flush_batch()
                             return
@@ -153,7 +129,7 @@ class QianbiDownloader(BaseDownloader):
 
                     # Normal chapter
                     batch.append(item)
-                    if len(batch) >= self.storage_batch_size:
+                    if len(batch) >= self._storage_batch_size:
                         await flush_batch()
 
                     if cancelled():
@@ -170,15 +146,13 @@ class QianbiDownloader(BaseDownloader):
                         # Final flush of everything
                         await flush_batch()
                         # Wait for remaining STOPs so chapter workers can finish.
-                        while stop_count < self.workers:
+                        while stop_count < self._workers:
                             nxt = await save_q.get()
                             if isinstance(nxt, StopToken):
                                 stop_count += 1
                         return
 
             # --- stage: chapter worker ---
-            sem = asyncio.Semaphore(self.workers)
-
             async def chapter_worker() -> None:
                 """
                 Fetch + parse with retry, then enqueue to save_q.
@@ -192,25 +166,20 @@ class QianbiDownloader(BaseDownloader):
                         await save_q.put(STOP)
                         return
 
-                    if not cid or cid in ignore_set:
-                        # Ignore silently and continue.
-                        continue
-
                     # If cancelled, don't start a new network call; let storage finish.
                     if cancelled():
                         await save_q.put(STOP)
                         return
 
-                    async with sem:
-                        chap = await self._process_chapter(book_id, cid, html_dir)
+                    chap = await self._process_chapter(book_id, cid, html_dir)
                     if chap:
                         await save_q.put(chap)
 
                     # polite pacing
                     await async_jitter_sleep(
-                        self.request_interval,
+                        self._request_interval,
                         mul_spread=1.1,
-                        max_sleep=self.request_interval + 2,
+                        max_sleep=self._request_interval + 2,
                     )
 
             # --- stage: producer ---
@@ -221,22 +190,22 @@ class QianbiDownloader(BaseDownloader):
                 so chapter workers can exit deterministically.
                 """
                 try:
-                    async for cid in self._chapter_ids(vols, start_id, end_id):
+                    for cid in plan:
                         if cancelled():
                             break
-                        if self.skip_existing and chapter_storage.exists(cid):
+                        if self._skip_existing and storage.exists(cid):
                             # Count as completed but don't enqueue.
                             await progress.bump(1)
                         else:
                             await cid_q.put(cid)
                 finally:
-                    for _ in range(self.workers):
+                    for _ in range(self._workers):
                         await cid_q.put(STOP)
 
             # --- run the pipeline ---
             async with asyncio.TaskGroup() as tg:
                 tg.create_task(storage_worker())
-                for _ in range(self.workers):
+                for _ in range(self._workers):
                     tg.create_task(chapter_worker())
                 tg.create_task(producer())
 
@@ -256,9 +225,6 @@ class QianbiDownloader(BaseDownloader):
                     book_info.get("book_name", "unknown"),
                 )
 
-        finally:
-            chapter_storage.close()
-
     async def _repair_chapter_ids(
         self,
         book_id: str,
@@ -295,9 +261,9 @@ class QianbiDownloader(BaseDownloader):
                         continue
                     storage.upsert_chapter(data, self.DEFAULT_SOURCE_ID)
                     await async_jitter_sleep(
-                        self.request_interval,
+                        self._request_interval,
                         mul_spread=1.1,
-                        max_sleep=self.request_interval + 2,
+                        max_sleep=self._request_interval + 2,
                     )
 
                 next_cid = data.get("extra", {}).get("next_chapter_id")
@@ -331,7 +297,7 @@ class QianbiDownloader(BaseDownloader):
 
         :return: ChapterDict on success, or None on failure.
         """
-        for attempt in range(self.retry_times + 1):
+        for attempt in range(self._retry_times + 1):
             try:
                 html_list = await self.fetcher.get_book_chapter(book_id, cid)
                 self._save_html_pages(html_dir, cid, html_list)
@@ -342,9 +308,9 @@ class QianbiDownloader(BaseDownloader):
                     raise ValueError("Empty parse result")
                 return chap
             except Exception as e:
-                if attempt < self.retry_times:
+                if attempt < self._retry_times:
                     self.logger.info(f"[ChapterWorker] Retry {cid} ({attempt+1}): {e}")
-                    backoff = self.backoff_factor * (2**attempt)
+                    backoff = self._backoff_factor * (2**attempt)
                     await async_jitter_sleep(
                         base=backoff, mul_spread=1.2, max_sleep=backoff + 3
                     )