refactor: optimize structure, stability and runtime performance

routes/api_screenshots.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import os
 from datetime import datetime
+from typing import Iterator
 
 import database
 from app_config import get_config
@@ -15,41 +16,67 @@ from services.time_utils import BEIJING_TZ
 
 config = get_config()
 SCREENSHOTS_DIR = config.SCREENSHOTS_DIR
+_IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg")
 
 api_screenshots_bp = Blueprint("api_screenshots", __name__)
 
 
+def _get_user_prefix(user_id: int) -> str:
+    user_info = database.get_user_by_id(user_id)
+    return user_info["username"] if user_info else f"user{user_id}"
+
+
+def _is_user_screenshot(filename: str, username_prefix: str) -> bool:
+    return filename.startswith(username_prefix + "_") and filename.lower().endswith(_IMAGE_EXTENSIONS)
+
+
+def _iter_user_screenshot_entries(username_prefix: str) -> Iterator[os.DirEntry]:
+    if not os.path.exists(SCREENSHOTS_DIR):
+        return
+
+    with os.scandir(SCREENSHOTS_DIR) as entries:
+        for entry in entries:
+            if (not entry.is_file()) or (not _is_user_screenshot(entry.name, username_prefix)):
+                continue
+            yield entry
+
+
+def _build_display_name(filename: str) -> str:
+    base_name, ext = filename.rsplit(".", 1)
+    parts = base_name.split("_", 1)
+    if len(parts) > 1:
+        return f"{parts[1]}.{ext}"
+    return filename
+
+
 @api_screenshots_bp.route("/api/screenshots", methods=["GET"])
 @login_required
 def get_screenshots():
     """获取当前用户的截图列表"""
     user_id = current_user.id
-    user_info = database.get_user_by_id(user_id)
-    username_prefix = user_info["username"] if user_info else f"user{user_id}"
+    username_prefix = _get_user_prefix(user_id)
 
     try:
         screenshots = []
-        if os.path.exists(SCREENSHOTS_DIR):
-            for filename in os.listdir(SCREENSHOTS_DIR):
-                if filename.lower().endswith((".png", ".jpg", ".jpeg")) and filename.startswith(username_prefix + "_"):
-                    filepath = os.path.join(SCREENSHOTS_DIR, filename)
-                    stat = os.stat(filepath)
-                    created_time = datetime.fromtimestamp(stat.st_mtime, tz=BEIJING_TZ)
-                    parts = filename.rsplit(".", 1)[0].split("_", 1)
-                    if len(parts) > 1:
-                        display_name = parts[1] + "." + filename.rsplit(".", 1)[1]
-                    else:
-                        display_name = filename
+        for entry in _iter_user_screenshot_entries(username_prefix):
+            filename = entry.name
+            stat = entry.stat()
+            created_time = datetime.fromtimestamp(stat.st_mtime, tz=BEIJING_TZ)
+
+            screenshots.append(
+                {
+                    "filename": filename,
+                    "display_name": _build_display_name(filename),
+                    "size": stat.st_size,
+                    "created": created_time.strftime("%Y-%m-%d %H:%M:%S"),
+                    "_created_ts": stat.st_mtime,
+                }
+            )
+
+        screenshots.sort(key=lambda item: item.get("_created_ts", 0), reverse=True)
+        for item in screenshots:
+            item.pop("_created_ts", None)
 
-                    screenshots.append(
-                        {
-                            "filename": filename,
-                            "display_name": display_name,
-                            "size": stat.st_size,
-                            "created": created_time.strftime("%Y-%m-%d %H:%M:%S"),
-                        }
-                    )
-        screenshots.sort(key=lambda x: x["created"], reverse=True)
         return jsonify(screenshots)
     except Exception as e:
         return jsonify({"error": str(e)}), 500
@@ -60,10 +87,9 @@ def get_screenshots():
 def serve_screenshot(filename):
     """提供截图文件访问"""
     user_id = current_user.id
-    user_info = database.get_user_by_id(user_id)
-    username_prefix = user_info["username"] if user_info else f"user{user_id}"
+    username_prefix = _get_user_prefix(user_id)
 
-    if not filename.startswith(username_prefix + "_"):
+    if not _is_user_screenshot(filename, username_prefix):
         return jsonify({"error": "无权访问"}), 403
 
     if not is_safe_path(SCREENSHOTS_DIR, filename):
@@ -77,12 +103,14 @@ def serve_screenshot(filename):
 def delete_screenshot(filename):
     """删除指定截图"""
     user_id = current_user.id
-    user_info = database.get_user_by_id(user_id)
-    username_prefix = user_info["username"] if user_info else f"user{user_id}"
+    username_prefix = _get_user_prefix(user_id)
 
-    if not filename.startswith(username_prefix + "_"):
+    if not _is_user_screenshot(filename, username_prefix):
         return jsonify({"error": "无权删除"}), 403
 
+    if not is_safe_path(SCREENSHOTS_DIR, filename):
+        return jsonify({"error": "非法路径"}), 403
+
     try:
         filepath = os.path.join(SCREENSHOTS_DIR, filename)
         if os.path.exists(filepath):
@@ -99,19 +127,15 @@ def delete_screenshot(filename):
 def clear_all_screenshots():
     """清空当前用户的所有截图"""
     user_id = current_user.id
-    user_info = database.get_user_by_id(user_id)
-    username_prefix = user_info["username"] if user_info else f"user{user_id}"
+    username_prefix = _get_user_prefix(user_id)
 
     try:
         deleted_count = 0
-        if os.path.exists(SCREENSHOTS_DIR):
-            for filename in os.listdir(SCREENSHOTS_DIR):
-                if filename.lower().endswith((".png", ".jpg", ".jpeg")) and filename.startswith(username_prefix + "_"):
-                    filepath = os.path.join(SCREENSHOTS_DIR, filename)
-                    os.remove(filepath)
-                    deleted_count += 1
+        for entry in _iter_user_screenshot_entries(username_prefix):
+            os.remove(entry.path)
+            deleted_count += 1
+
         log_to_client(f"清理了 {deleted_count} 个截图文件", user_id)
         return jsonify({"success": True, "deleted": deleted_count})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
-