From 3af48126eafd29d351c976ac79b9ab76c23b4445 Mon Sep 17 00:00:00 2001
From: g1879 <g1879@qq.com>
Date: Fri, 7 Jan 2022 18:15:20 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E5=85=B3=E9=97=AD=E6=A0=87?=
 =?UTF-8?q?=E7=AD=BE=E9=A1=B5=E5=8A=9F=E8=83=BD=EF=BC=8C=E6=9C=AA=E5=AE=8C?=
 =?UTF-8?q?=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 DrissionPage/driver_page.py | 43 +++++++++++++++++++++++--------------
 1 file changed, 27 insertions(+), 16 deletions(-)

diff --git a/DrissionPage/driver_page.py b/DrissionPage/driver_page.py
index 2b345c4..f8941d8 100644
--- a/DrissionPage/driver_page.py
+++ b/DrissionPage/driver_page.py
@@ -303,31 +303,28 @@ class DriverPage(BasePage):
         :param num_or_handles: 要保留的标签页序号或handle，可传入handle组成的列表或元组
         :return: None
         """
-        try:
-            tab = int(num_or_handles)
-        except (ValueError, TypeError):
-            tab = num_or_handles
-
-        tabs = self.driver.window_handles
-
-        if tab is None:
-            page_handle = (self.current_tab_handle,)
-        elif isinstance(tab, int):
-            page_handle = (tabs[tab],)
-        elif isinstance(tab, str):
-            page_handle = (tab,)
-        elif isinstance(tab, (list, tuple)):
-            page_handle = tab
+        if num_or_handles is None:
+            reserve_tabs = (self.current_tab_handle,)
+        elif isinstance(num_or_handles, (int, str)):
+            reserve_tabs = (num_or_handles,)
+        elif isinstance(num_or_handles, (list, tuple)):
+            reserve_tabs = num_or_handles
         else:
             raise TypeError('num_or_handle参数只能是int、str、list 或 tuple类型。')
 
+        tabs = self.driver.window_handles
+        reserve_tabs = set(i if isinstance(i, str) else tabs[i] for i in reserve_tabs)
+
         for i in tabs:  # 遍历所有标签页，关闭非保留的
-            if i not in page_handle:
+            if i not in reserve_tabs:
                 self.driver.switch_to.window(i)
                 self.driver.close()
 
         self.to_tab(0)
 
+    def close_tab(self, num_or_handles: Union[int, str, list, tuple] = None):
+        pass
+
     def to_tab(self, num_or_handle: Union[int, str] = 0) -> None:
         """跳转到标签页                                                         \n
         :param num_or_handle: 标签页序号或handle字符串，序号第一个为0，最后为-1
@@ -525,3 +522,17 @@ class DriverPage(BasePage):
             alert.accept()
 
         return text
+
+
+def _get_handles(page: DriverPage, num_or_handles: Union[int, str, list, tuple] = None):
+    if num_or_handles is None:
+        reserve_tabs = (page.current_tab_handle,)
+    elif isinstance(num_or_handles, (int, str)):
+        reserve_tabs = (num_or_handles,)
+    elif isinstance(num_or_handles, (list, tuple)):
+        reserve_tabs = num_or_handles
+    else:
+        raise TypeError('num_or_handle参数只能是int、str、list 或 tuple类型。')
+
+    tabs = page.driver.window_handles
+    return set(i if isinstance(i, str) else tabs[i] for i in reserve_tabs)