From 6a063787ee83a3b67b68b912ca1ff90819453de8 Mon Sep 17 00:00:00 2001
From: g1879 <g1879@qq.com>
Date: Tue, 12 Dec 2023 21:24:08 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0set.auto=5Fhandle=5Falert()?=
 =?UTF-8?q?=EF=BC=9Bhandle=5Falert()=E5=A2=9E=E5=8A=A0next=5Fone=E5=8F=82?=
 =?UTF-8?q?=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 DrissionPage/_pages/chromium_base.py  | 37 +++++++++++++++++++--------
 DrissionPage/_pages/chromium_base.pyi |  6 ++++-
 DrissionPage/_units/setter.py         |  8 ++++++
 DrissionPage/_units/setter.pyi        |  2 ++
 4 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/DrissionPage/_pages/chromium_base.py b/DrissionPage/_pages/chromium_base.py
index eb31600..a48db70 100644
--- a/DrissionPage/_pages/chromium_base.py
+++ b/DrissionPage/_pages/chromium_base.py
@@ -835,13 +835,19 @@ class ChromiumBase(BasePage):
         if cookies:
             self.run_cdp_loaded('Network.clearBrowserCookies')
 
-    def handle_alert(self, accept=True, send=None, timeout=None):
+    def handle_alert(self, accept=True, send=None, timeout=None, next_one=False):
         """处理提示框，可以自动等待提示框出现
         :param accept: True表示确认，False表示取消，其它值不会按按钮但依然返回文本值
         :param send: 处理prompt提示框时可输入文本
         :param timeout: 等待提示框出现的超时时间，为None则使用self.timeout属性的值
+        :param next_one: 是否处理下一个出现的提示框，为True时timeout参数无效
         :return: 提示框内容文本，未等到提示框则返回False
         """
+        if next_one:
+            self._alert.handle_next = accept
+            self._alert.next_text = send
+            return
+
         timeout = self.timeout if timeout is None else timeout
         timeout = .1 if timeout <= 0 else timeout
         end_time = perf_counter() + timeout
@@ -857,16 +863,6 @@ class ChromiumBase(BasePage):
             self.driver.run('Page.handleJavaScriptDialog', accept=accept)
         return res_text
 
-    def _on_alert_close(self, **kwargs):
-        """alert关闭时触发的方法"""
-        self._alert.activated = False
-        self._alert.text = None
-        self._alert.type = None
-        self._alert.defaultPrompt = None
-        self._alert.response_accept = kwargs.get('result')
-        self._alert.response_text = kwargs['userInput']
-        self._has_alert = False
-
     def _on_alert_open(self, **kwargs):
         """alert出现时触发的方法"""
         self._alert.activated = True
@@ -877,6 +873,22 @@ class ChromiumBase(BasePage):
         self._alert.response_text = None
         self._has_alert = True
 
+        if self._alert.auto is not None:
+            self.handle_alert(self._alert.auto)
+        elif self._alert.handle_next is not None:
+            self.handle_alert(self._alert.handle_next, self._alert.next_text)
+            self._alert.handle_next = None
+
+    def _on_alert_close(self, **kwargs):
+        """alert关闭时触发的方法"""
+        self._alert.activated = False
+        self._alert.text = None
+        self._alert.type = None
+        self._alert.defaultPrompt = None
+        self._alert.response_accept = kwargs.get('result')
+        self._alert.response_text = kwargs['userInput']
+        self._has_alert = False
+
     def _wait_loaded(self, timeout=None):
         """等待页面加载完成，超时触发停止加载
         :param timeout: 超时时间
@@ -1107,6 +1119,9 @@ class Alert(object):
         self.defaultPrompt = None
         self.response_accept = None
         self.response_text = None
+        self.handle_next = None
+        self.next_text = None
+        self.auto = None
 
 
 def close_privacy_dialog(page, tid):
diff --git a/DrissionPage/_pages/chromium_base.pyi b/DrissionPage/_pages/chromium_base.pyi
index 24c9154..3624a20 100644
--- a/DrissionPage/_pages/chromium_base.pyi
+++ b/DrissionPage/_pages/chromium_base.pyi
@@ -227,7 +227,8 @@ class ChromiumBase(BasePage):
     def clear_cache(self, session_storage: bool = True, local_storage: bool = True, cache: bool = True,
                     cookies: bool = True) -> None: ...
 
-    def handle_alert(self, accept: bool = True, send: str = None, timeout: float = None) -> Union[str, False]: ...
+    def handle_alert(self, accept: bool = True, send: str = None, timeout: float = None,
+                     next_one: bool = False) -> Union[str, False]: ...
 
     def _on_alert_close(self, **kwargs): ...
 
@@ -257,3 +258,6 @@ class Alert(object):
         self.defaultPrompt: str = ...
         self.response_accept: str = ...
         self.response_text: str = ...
+        self.handle_next: Optional[bool] = ...
+        self.next_text: str = ...
+        self.auto: Optional[bool] = ...
diff --git a/DrissionPage/_units/setter.py b/DrissionPage/_units/setter.py
index cfca821..4b5b952 100644
--- a/DrissionPage/_units/setter.py
+++ b/DrissionPage/_units/setter.py
@@ -146,6 +146,14 @@ class ChromiumBaseSetter(BasePageSetter):
         self._page.run_cdp('Network.enable')
         self._page.run_cdp('Network.setExtraHTTPHeaders', headers=headers)
 
+    def auto_handle_alert(self, on_off=True, accept=True):
+        """设置是否启用自动处理弹窗
+        :param on_off: bool表示开或关
+        :param accept: bool表示确定还是取消
+        :return: None
+        """
+        self._page._alert.auto = accept if on_off else None
+
     # --------------即将废弃---------------
 
     @property
diff --git a/DrissionPage/_units/setter.pyi b/DrissionPage/_units/setter.pyi
index 45faf59..19e2196 100644
--- a/DrissionPage/_units/setter.pyi
+++ b/DrissionPage/_units/setter.pyi
@@ -59,6 +59,8 @@ class ChromiumBaseSetter(BasePageSetter):
 
     def headers(self, headers: dict) -> None: ...
 
+    def auto_handle_alert(self, on_off: bool = True, accept: bool = True) -> None: ...
+
     def upload_files(self, files: Union[str, list, tuple]) -> None: ...