Improve KDocs QR capture

services/kdocs_uploader.py

@@ -268,24 +268,104 @@ class KDocsUploader:
             pass
 
     def _capture_qr_image(self) -> Optional[bytes]:
-        candidates = [
-            self._page.locator("canvas"),
-            self._page.locator("img[alt*='二维码']"),
-            self._page.locator("img[src*='qr']"),
-        ]
-        for locator in candidates:
-            try:
-                if locator.count() < 1:
+        pages = self._iter_pages()
+        for page in pages:
+            for frame in page.frames:
+                target = self._find_qr_element_in_frame(frame)
+                if not target:
                     continue
-                target = locator.first
-                target.wait_for(state="visible", timeout=5000)
-                return target.screenshot()
-            except PlaywrightTimeoutError:
-                continue
+                try:
+                    return target.screenshot()
+                except Exception:
+                    continue
+        for page in pages:
+            dialog_image = self._capture_dialog_image(page)
+            if dialog_image:
+                return dialog_image
+        return None
+
+    def _iter_pages(self) -> list:
+        pages = []
+        if self._context:
+            pages.extend(self._context.pages)
+        if self._page and self._page not in pages:
+            pages.insert(0, self._page)
+        def rank(p) -> int:
+            url = (getattr(p, "url", "") or "").lower()
+            keywords = ("login", "account", "passport", "wechat", "qr")
+            return 0 if any(k in url for k in keywords) else 1
+        pages.sort(key=rank)
+        return pages
+
+    def _find_qr_element_in_frame(self, frame) -> Optional[Any]:
+        selectors = [
+            "canvas",
+            "img[alt*='二维码']",
+            "img[src*='qr']",
+            "img[src*='qrcode']",
+            "img[class*='qr']",
+            "canvas[class*='qr']",
+            "svg",
+            "div[role='img']",
+        ]
+        best = None
+        best_score = None
+        for selector in selectors:
+            try:
+                locator = frame.locator(selector)
+                count = min(locator.count(), 20)
             except Exception:
                 continue
+            for i in range(count):
+                el = locator.nth(i)
+                try:
+                    if not el.is_visible(timeout=800):
+                        continue
+                    box = el.bounding_box()
+                    if not box:
+                        continue
+                    width = box.get("width", 0)
+                    height = box.get("height", 0)
+                    if width < 120 or height < 120 or width > 420 or height > 420:
+                        continue
+                    aspect_diff = abs(width - height)
+                    if aspect_diff > 40:
+                        continue
+                    score = aspect_diff + abs(width - 260) + abs(height - 260)
+                    if best_score is None or score < best_score:
+                        best_score = score
+                        best = el
+                except Exception:
+                    continue
+        return best
+
+    def _capture_dialog_image(self, page) -> Optional[bytes]:
+        selectors = "[role='dialog'], .dialog, .modal, .popup"
         try:
-            return self._page.screenshot(full_page=True)
+            dialogs = page.locator(selectors)
+            count = min(dialogs.count(), 6)
+        except Exception:
+            return None
+        best = None
+        best_area = 0
+        for i in range(count):
+            el = dialogs.nth(i)
+            try:
+                if not el.is_visible(timeout=800):
+                    continue
+                box = el.bounding_box()
+                if not box:
+                    continue
+                area = box.get("width", 0) * box.get("height", 0)
+                if area > best_area:
+                    best_area = area
+                    best = el
+            except Exception:
+                continue
+        if not best:
+            return None
+        try:
+            return best.screenshot()
         except Exception:
             return None
 
@@ -317,7 +397,12 @@ class KDocsUploader:
             return {"success": True, "logged_in": True, "qr_image": ""}
 
         self._ensure_login_dialog()
-        qr_image = self._capture_qr_image()
+        qr_image = None
+        for _ in range(5):
+            qr_image = self._capture_qr_image()
+            if qr_image:
+                break
+            time.sleep(1)
         if not qr_image:
             self._last_error = "二维码获取失败"
             return {"success": False, "error": self._last_error}