diff --git a/app/services/book_service.py b/app/services/book_service.py index 92aee83..3960223 100644 --- a/app/services/book_service.py +++ b/app/services/book_service.py @@ -1,10 +1,9 @@ -import os import requests from pathlib import Path from bs4 import BeautifulSoup from app.extensions import db from tenacity import retry, stop_after_attempt, RetryError -from app.utils import datetime_to_timestamp, timestamp_to_datetime, get_today_timestamp +from app.utils import datetime_to_timestamp, timestamp_to_datetime, get_today_timestamp, find_project_root from app.models import Setting @@ -22,8 +21,9 @@ class BookService: target_ts = datetime_to_timestamp(self.setting["targetDate"]) - 24 * 60 * 60 self.setting["targetDate"] = timestamp_to_datetime(target_ts) - app_dir = os.path.dirname(os.path.dirname(__file__)) - self.download_path = os.path.join(os.path.dirname(app_dir), "download") # 下载目录 + app_dir = find_project_root() + print(app_dir) + self.download_path = app_dir.parent / "download" # 下载目录 self.today_timestamp = get_today_timestamp() self.one_date_loading = False self.session = requests.session() @@ -54,16 +54,11 @@ class BookService: name = name.replace(old, new) return name - def create_folder(self, target_dir, name): # 创建文件夹 - real_dir = os.path.join(target_dir, name) - folder = os.path.exists(real_dir) - if not folder: - os.makedirs(real_dir) - else: - print(f"文件夹：{name} 已存在") - - def isDownloadTarget(self, soup): # 是否是可以下载的文件 - return soup.headers.get('Content-Type') in self.ALLOWED_CONTENT_TYPES + def create_folder(self, parent_dir: Path, folder_name: str) -> Path: # 创建文件夹 + """创建文件夹（如果不存在）并返回路径""" + target = parent_dir / folder_name + target.mkdir(parents=True, exist_ok=True) + return target def post_page(self, name, url): # 帖子页面 print(f"详情页面 {self.setting["baseUrl"]}{url}") @@ -73,7 +68,7 @@ class BookService: soup = BeautifulSoup(response.text, 'lxml') resource_boxs = soup.select('ignore_js_op') for i in resource_boxs: - download_dir_path = os.path.join(self.download_path, self.setting["targetDate"]) + download_dir_path = self.download_path / self.setting["targetDate"] self.create_folder(download_dir_path, name) try: for string in i.select_one(".tip.tip_4").stripped_strings: @@ -89,20 +84,26 @@ class BookService: print('重试三次也不成功') @retry(stop=stop_after_attempt(3)) - def download_file(self, file_url, dir_name, file_name): # 下载文件 - response = self.session.get(file_url, stream=True) - response.raise_for_status() - print(f"下载状态码：{response.status_code} {file_url}") - print(dir_name + "/" + file_name) - if self.isDownloadTarget(response): - with open(f"download/{self.setting["targetDate"]}/{dir_name}/{file_name}", 'wb') as file2: + def download_file(self, file_url: str, dir_name: str, file_name: str): # 下载文件 + try: + response = self.session.get(file_url, stream=True) + response.raise_for_status() + print(f"下载状态码：{response.status_code} {file_url}") + print(dir_name + "/" + file_name) + content_type = response.headers.get('Content-Type', '') + if content_type not in self.ALLOWED_CONTENT_TYPES: + yield f"event: mes_error\ndata: {self.setting["targetDate"]} {file_name}\n\n" + print("不是可下载文件") + return + target_dir = f"download/{self.setting["targetDate"]}/{dir_name}/{file_name}" + with open(target_dir, 'wb') as f: for chunk in response.iter_content(chunk_size=1024): if chunk: - file2.write(chunk) + f.write(chunk) yield f"data: {self.setting["targetDate"]} {file_name}\n\n" - else: - yield f"event: mes_error\ndata: {self.setting["targetDate"]} {file_name}\n\n" - print("不是可下载文件") + except requests.RequestException as e: + yield f"event: error\ndata: {self.setting['targetDate']} {file_name} 网络错误\n\n" + raise # 触发 tenacity 重试 def section_page(self, page): # 版块页面 print(f"当前页面 {page}") @@ -169,12 +170,17 @@ class BookService: yield from self.section_page(self.setting["totalPage"] - self.setting["targetCountdownPage"]) return None - def get_post_date(self, element): - if element.select_one(".by em span span"): - if element.select_one(".by em span span").get("title"): - return element.select_one(".by em span span").get("title") + # ------------------------------------------------------------ + # 获取帖子发布时间 + # ------------------------------------------------------------ + @staticmethod + def get_post_date(element): + span_span = element.select_one(".by em span span") + if span_span: + if span_span.get("title"): + return span_span.get("title") else: - return element.select_one(".by em span span").string + return span_span.string else: return element.select_one(".by em span").string diff --git a/app/utils/__init__.py b/app/utils/__init__.py index 4a0d484..745f987 100644 --- a/app/utils/__init__.py +++ b/app/utils/__init__.py @@ -2,6 +2,7 @@ import json import time import ctypes from collections import defaultdict +from pathlib import Path from typing import List from datetime import datetime from functools import cmp_to_key @@ -103,4 +104,11 @@ def to_json_serializable(obj): if isinstance(obj, set): return list(obj) # 集合转列表 # 其他基本类型（str, int, float, list, tuple, None）直接返回 - return obj \ No newline at end of file + return obj + +def find_project_root(marker_files=('requirements.txt', 'run.py')): + current = Path(__file__).resolve().parent + for parent in current.parents: + if any((parent / marker).exists() for marker in marker_files): + return parent + return current # 没找到则返回当前文件所在目录 \ No newline at end of file