diff --git a/app.py b/app.py index 3b2414d1..b7eb94e7 100644 --- a/app.py +++ b/app.py @@ -796,7 +796,7 @@ def handle_connect(): @socketio.on('disconnect') def handle_disconnect(): - """Handle client disconnection.""" + """Handle client disconnection.""" logger.info("WebSocket client disconnected") @socketio.on('request_status') diff --git a/backend.py b/backend.py index 868bd67e..3fab4699 100644 --- a/backend.py +++ b/backend.py @@ -158,6 +158,9 @@ def _download_book_with_cancellation(book_id: str, cancel_flag: Event) -> Option book_name = _sanitize_filename(book_info.title) else: book_name = book_id + # If format is not set, use the format of the first download URL + if book_info.format == "": + book_info.format = book_info.download_urls[0].split(".")[-1] book_name += f".{book_info.format}" book_path = TMP_DIR / book_name @@ -168,7 +171,7 @@ def _download_book_with_cancellation(book_id: str, cancel_flag: Event) -> Option progress_callback = lambda progress: update_download_progress(book_id, progress) status_callback = lambda status: update_download_status(book_id, status) - success = book_manager.download_book(book_info, book_path, progress_callback, cancel_flag, status_callback) + success_download_url = book_manager.download_book(book_info, book_path, progress_callback, cancel_flag, status_callback) # Stop progress updates cancel_flag.wait(0.1) # Brief pause for progress thread cleanup @@ -180,7 +183,7 @@ def _download_book_with_cancellation(book_id: str, cancel_flag: Event) -> Option book_path.unlink() return None - if not success: + if not success_download_url: raise Exception("Unknown error downloading book") # Check cancellation before post-processing @@ -205,6 +208,10 @@ def _download_book_with_cancellation(book_id: str, cancel_flag: Event) -> Option if ws_manager and ws_manager.is_enabled(): ws_manager.broadcast_status_update(queue_status()) + if success_download_url and book_info.format == "": + book_info.format = success_download_url.split(".")[-1] + book_name += f".{book_info.format}" + intermediate_path = INGEST_DIR / f"{book_id}.crdownload" final_path = INGEST_DIR / book_name diff --git a/book_manager.py b/book_manager.py index f6198b08..77111ca3 100644 --- a/book_manager.py +++ b/book_manager.py @@ -229,22 +229,24 @@ def _parse_book_info_page(soup: BeautifulSoup, book_id: str) -> BookInfo: original_divs = divs divs = [div for div in divs if div.text.strip() != ""] - _details = _find_in_divs(divs, " · ").split(" · ") + all_details = _find_in_divs(divs, " · ") format = "" size = "" - for f in _details: - if format == "" and f.strip().lower() in SUPPORTED_FORMATS: - format = f.strip().lower() - if size == "" and any(u in f.strip().lower() for u in ["mb", "kb", "gb"]): - size = f.strip().lower() - - if format == "" or size == "": + for _details in all_details: + _details = _details.split(" · ") for f in _details: - stripped = f.strip().lower() - if format == "" and stripped and " " not in stripped: - format = stripped - if size == "" and "." in stripped: - size = stripped + if format == "" and f.strip().lower() in SUPPORTED_FORMATS: + format = f.strip().lower() + if size == "" and any(u in f.strip().lower() for u in ["mb", "kb", "gb"]): + size = f.strip().lower() + + if format == "" or size == "": + for f in _details: + stripped = f.strip().lower() + if format == "" and stripped and " " not in stripped: + format = stripped + if size == "" and "." in stripped: + size = stripped book_title = _find_in_divs(divs, "🔍").strip("🔍").strip() @@ -280,15 +282,16 @@ def _parse_book_info_page(soup: BeautifulSoup, book_id: str) -> BookInfo: return book_info -def _find_in_divs(divs: List[str], text: str, isClass: bool = False) -> str: +def _find_in_divs(divs: List[str], text: str, isClass: bool = False) -> List[str]: + divs_found = [] for div in divs: if isClass: if div.find(class_ = text): - return div.text.strip() + divs_found.append(div.text.strip()) else: if text in div.text.strip(): - return div.text.strip() - return "" + divs_found.append(div.text.strip()) + return divs_found def _get_download_urls_from_welib(book_id: str) -> set[str]: if ALLOW_USE_WELIB == False: @@ -390,7 +393,7 @@ def _extract_book_metadata(metadata_divs) -> Dict[str, List[str]]: } -def download_book(book_info: BookInfo, book_path: Path, progress_callback: Optional[Callable[[float], None]] = None, cancel_flag: Optional[Event] = None, status_callback: Optional[Callable[[str], None]] = None) -> bool: +def download_book(book_info: BookInfo, book_path: Path, progress_callback: Optional[Callable[[float], None]] = None, cancel_flag: Optional[Event] = None, status_callback: Optional[Callable[[str], None]] = None) -> Optional[str]: """Download a book from available sources. Args: @@ -401,7 +404,7 @@ def download_book(book_info: BookInfo, book_path: Path, progress_callback: Optio status_callback: Optional callback for status updates Returns: - bool: True if successful, False otherwise + str: Download URL if successful, None otherwise """ if len(book_info.download_urls) == 0: @@ -437,13 +440,13 @@ def download_book(book_info: BookInfo, book_path: Path, progress_callback: Optio with open(book_path, "wb") as f: f.write(data.getbuffer()) logger.info(f"Writing `{book_info.title}` successfully") - return True + return download_url except Exception as e: logger.error_trace(f"Failed to download from {link}: {e}") continue - return False + return None def _get_download_url(link: str, title: str, cancel_flag: Optional[Event] = None, status_callback: Optional[Callable[[str], None]] = None) -> str: