Speed up KDocs QR retrieval

services/kdocs_uploader.py

@@ -199,11 +199,31 @@ class KDocsUploader:
             pass
         self._playwright = None
 
-    def _open_document(self, doc_url: str) -> bool:
+    def _open_document(self, doc_url: str, *, fast: bool = False) -> bool:
         try:
             self._doc_url = doc_url
-            self._page.goto(doc_url)
-            time.sleep(1)
+            if fast:
+                doc_pages = self._find_doc_pages(doc_url)
+                if doc_pages:
+                    self._page = doc_pages[0]
+                    return True
+                login_pages = []
+                for page in self._list_pages():
+                    try:
+                        url = getattr(page, "url", "") or ""
+                        if self._is_login_url(url):
+                            login_pages.append(page)
+                    except Exception:
+                        continue
+                if login_pages:
+                    self._page = login_pages[0]
+                    return True
+            goto_kwargs = {}
+            if fast:
+                fast_timeout = int(os.environ.get("KDOCS_FAST_GOTO_TIMEOUT_MS", "15000"))
+                goto_kwargs = {"wait_until": "domcontentloaded", "timeout": fast_timeout}
+            self._page.goto(doc_url, **goto_kwargs)
+            time.sleep(0.6)
             doc_pages = self._find_doc_pages(doc_url)
             if doc_pages and doc_pages[0] is not self._page:
                 self._page = doc_pages[0]
@@ -315,22 +335,40 @@ class KDocsUploader:
         storage_state = getattr(config, "KDOCS_LOGIN_STATE_FILE", "data/kdocs_login_state.json")
         return os.path.exists(storage_state)
 
-    def _ensure_login_dialog(self) -> None:
+    def _ensure_login_dialog(
+        self,
+        *,
+        timeout_ms: int = 1200,
+        frame_timeout_ms: int = 800,
+        quick: bool = False,
+    ) -> None:
         login_names = ["登录并加入编辑", "立即登录", "登录"]
         wechat_names = ["微信登录", "微信扫码登录", "微信扫码", "扫码登录"]
         pages = self._iter_pages()
         clicked = False
         for page in pages:
-            if self._try_click_names(page, login_names):
+            if self._try_click_names(
+                page,
+                login_names,
+                timeout_ms=timeout_ms,
+                frame_timeout_ms=frame_timeout_ms,
+                quick=quick,
+            ):
                 clicked = True
                 break
         if clicked:
             time.sleep(1.5)
         pages = self._iter_pages()
         for page in pages:
-            if self._try_click_names(page, wechat_names):
+            if self._try_click_names(
+                page,
+                wechat_names,
+                timeout_ms=timeout_ms,
+                frame_timeout_ms=frame_timeout_ms,
+                quick=quick,
+            ):
                 return
-        self._try_confirm_login()
+        self._try_confirm_login(timeout_ms=timeout_ms, frame_timeout_ms=frame_timeout_ms, quick=quick)
 
     def _capture_qr_image(self) -> Optional[bytes]:
         pages = self._iter_pages()
@@ -458,34 +496,44 @@ class KDocsUploader:
             return False
         return False
 
-    def _try_click_names(self, page, names: list) -> bool:
+    def _try_click_names(
+        self,
+        page,
+        names: list,
+        *,
+        timeout_ms: int = 1200,
+        frame_timeout_ms: int = 800,
+        quick: bool = False,
+    ) -> bool:
         for name in names:
-            if self._try_click_role(page, "button", name, timeout=1200):
-                return True
-            if self._try_click_role(page, "link", name, timeout=1200):
+            if self._try_click_role(page, "button", name, timeout=timeout_ms):
                 return True
+            if not quick:
+                if self._try_click_role(page, "link", name, timeout=timeout_ms):
+                    return True
             try:
                 el = page.get_by_text(name, exact=True)
-                if el.is_visible(timeout=1200):
-                    el.click()
-                    time.sleep(1)
-                    return True
-            except Exception:
-                pass
-            try:
-                el = page.get_by_text(name, exact=False)
-                if el.is_visible(timeout=1200):
+                if el.is_visible(timeout=timeout_ms):
                     el.click()
                     time.sleep(1)
                     return True
             except Exception:
                 pass
+            if not quick:
+                try:
+                    el = page.get_by_text(name, exact=False)
+                    if el.is_visible(timeout=timeout_ms):
+                        el.click()
+                        time.sleep(1)
+                        return True
+                except Exception:
+                    pass
         try:
             for frame in page.frames:
                 for name in names:
                     try:
                         el = frame.get_by_role("button", name=name)
-                        if el.is_visible(timeout=800):
+                        if el.is_visible(timeout=frame_timeout_ms):
                             el.click()
                             time.sleep(1)
                             return True
@@ -493,29 +541,42 @@ class KDocsUploader:
                         pass
                     try:
                         el = frame.get_by_text(name, exact=True)
-                        if el.is_visible(timeout=800):
-                            el.click()
-                            time.sleep(1)
-                            return True
-                    except Exception:
-                        pass
-                    try:
-                        el = frame.get_by_text(name, exact=False)
-                        if el.is_visible(timeout=800):
+                        if el.is_visible(timeout=frame_timeout_ms):
                             el.click()
                             time.sleep(1)
                             return True
                     except Exception:
                         pass
+                    if not quick:
+                        try:
+                            el = frame.get_by_text(name, exact=False)
+                            if el.is_visible(timeout=frame_timeout_ms):
+                                el.click()
+                                time.sleep(1)
+                                return True
+                        except Exception:
+                            pass
         except Exception:
             return False
         return False
 
-    def _try_confirm_login(self) -> bool:
+    def _try_confirm_login(
+        self,
+        *,
+        timeout_ms: int = 1200,
+        frame_timeout_ms: int = 800,
+        quick: bool = False,
+    ) -> bool:
         confirm_names = ["确认登录", "确认登陆"]
         pages = self._iter_pages()
         for page in pages:
-            if self._try_click_names(page, confirm_names):
+            if self._try_click_names(
+                page,
+                confirm_names,
+                timeout_ms=timeout_ms,
+                frame_timeout_ms=frame_timeout_ms,
+                quick=quick,
+            ):
                 return True
         return False
 
@@ -602,7 +663,7 @@ class KDocsUploader:
             self._handle_clear_login()
         if not self._ensure_playwright(use_storage_state=not force):
             return {"success": False, "error": self._last_error or "浏览器不可用"}
-        if not self._open_document(doc_url):
+        if not self._open_document(doc_url, fast=True):
             return {"success": False, "error": self._last_error or "打开文档失败"}
 
         if not force and self._has_saved_login_state() and self._is_logged_in():
@@ -611,11 +672,21 @@ class KDocsUploader:
             self._save_login_state()
             return {"success": True, "logged_in": True, "qr_image": ""}
 
-        self._ensure_login_dialog()
+        fast_login_timeout = int(os.environ.get("KDOCS_FAST_LOGIN_TIMEOUT_MS", "300"))
+        self._ensure_login_dialog(
+            timeout_ms=fast_login_timeout,
+            frame_timeout_ms=fast_login_timeout,
+            quick=True,
+        )
         qr_image = None
         invalid_qr = None
-        for _ in range(10):
-            self._ensure_login_dialog()
+        for attempt in range(10):
+            if attempt in (3, 7):
+                self._ensure_login_dialog(
+                    timeout_ms=fast_login_timeout,
+                    frame_timeout_ms=fast_login_timeout,
+                    quick=True,
+                )
             candidate = self._capture_qr_image()
             if candidate and self._is_valid_qr_image(candidate):
                 qr_image = candidate
@@ -689,10 +760,19 @@ class KDocsUploader:
             return {"success": True, "logged_in": False, "error": "未配置文档链接"}
         if not self._ensure_playwright():
             return {"success": False, "logged_in": False, "error": self._last_error or "浏览器不可用"}
-        if not self._open_document(doc_url):
+        if not self._open_document(doc_url, fast=True):
             return {"success": False, "logged_in": False, "error": self._last_error or "打开文档失败"}
-        self._ensure_login_dialog()
-        self._try_confirm_login()
+        fast_login_timeout = int(os.environ.get("KDOCS_FAST_LOGIN_TIMEOUT_MS", "300"))
+        self._ensure_login_dialog(
+            timeout_ms=fast_login_timeout,
+            frame_timeout_ms=fast_login_timeout,
+            quick=True,
+        )
+        self._try_confirm_login(
+            timeout_ms=fast_login_timeout,
+            frame_timeout_ms=fast_login_timeout,
+            quick=True,
+        )
         logged_in = self._is_logged_in()
         self._last_login_ok = logged_in
         self._login_required = not logged_in