From 7bb1e64a79cc73a870fb07bcff514b9181f65a57 Mon Sep 17 00:00:00 2001
From: g1879 <g1879@qq.com>
Date: Mon, 5 Feb 2024 15:47:55 +0800
Subject: [PATCH] =?UTF-8?q?4.0.4.5=E5=AE=8C=E5=96=84tree()?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 DrissionPage/_base/base.py      |  2 +-
 DrissionPage/_base/base.pyi     |  7 ++++--
 DrissionPage/_functions/web.py  | 44 +++++++++++++++++++--------------
 DrissionPage/_functions/web.pyi |  4 +--
 DrissionPage/common.py          |  4 +--
 5 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/DrissionPage/_base/base.py b/DrissionPage/_base/base.py
index 600a3f1..21e3292 100644
--- a/DrissionPage/_base/base.py
+++ b/DrissionPage/_base/base.py
@@ -36,7 +36,7 @@ class BaseParser(object):
     def html(self):
         return ''
 
-    def s_ele(self, locator):
+    def s_ele(self, locator=None):
         pass
 
     def s_eles(self, locator):
diff --git a/DrissionPage/_base/base.pyi b/DrissionPage/_base/base.pyi
index 2fabbe6..235d8f6 100644
--- a/DrissionPage/_base/base.pyi
+++ b/DrissionPage/_base/base.pyi
@@ -11,6 +11,7 @@ from typing import Union, Tuple, List, Any, Optional
 from DownloadKit import DownloadKit
 
 from .._elements.none_element import NoneElement
+from .._elements.session_element import SessionElement
 
 
 class BaseParser(object):
@@ -29,9 +30,11 @@ class BaseParser(object):
     @property
     def html(self) -> str: ...
 
-    def s_ele(self, locator: Union[Tuple[str, str], str, BaseElement], index: int = 1): ...
+    def s_ele(self,
+              locator: Union[Tuple[str, str], str, BaseElement, None] = None,
+              index: int = 1) -> SessionElement: ...
 
-    def s_eles(self, locator: Union[Tuple[str, str], str]): ...
+    def s_eles(self, locator: Union[Tuple[str, str], str]) -> List[SessionElement]: ...
 
     def _ele(self,
              locator: Union[Tuple[str, str], str],
diff --git a/DrissionPage/_functions/web.py b/DrissionPage/_functions/web.py
index 30387bf..1f22d09 100644
--- a/DrissionPage/_functions/web.py
+++ b/DrissionPage/_functions/web.py
@@ -365,24 +365,32 @@ def get_blob(page, url, as_bytes=True):
         return result
 
 
-def tree(ele_or_page, layer=5, last_one=False, body=''):
-    try:
-        list_ele = ele_or_page.s_ele().children(timeout=0.1)
-    except:
-        list_ele = []
-    length = len(list_ele)
-    body_unit = '    ' if last_one else '│   '
-    tail = '├───'
-    new_body = body + body_unit
+def tree(ele_or_page):
+    """把页面或元素对象DOM结构打印出来
+    :param ele_or_page: 页面或元素对象
+    :return: None
+    """
+    def _tree(obj, last_one=True, body=''):
+        list_ele = obj.children()
+        length = len(list_ele)
+        body_unit = '    ' if last_one else '│   '
+        tail = '├───'
+        new_body = body + body_unit
 
-    if length > 0 and layer >= 1:
-        new_last_one = False
-        for i in range(length):
-            if i == length - 1:
-                tail = '└───'
-                new_last_one = True
-            e = list_ele[i]
+        if length > 0:
+            new_last_one = False
+            for i in range(length):
+                if i == length - 1:
+                    tail = '└───'
+                    new_last_one = True
+                e = list_ele[i]
 
-            print(f'{new_body}{tail}{i}<{e.tag}>  {e.attrs}')
+                attrs = ' '.join([f"{k}='{v}'" for k, v in e.attrs.items()])
+                print(f'{new_body}{tail}<{e.tag} {attrs}>'.replace('\n', ' '))
 
-            tree(e, layer - 1, new_last_one, new_body)
\ No newline at end of file
+                _tree(e, new_last_one, new_body)
+
+    ele = ele_or_page.s_ele()
+    attrs = ' '.join([f"{k}='{v}'" for k, v in ele.attrs.items()])
+    print(f'<{ele.tag} {attrs}>'.replace('\n', ' '))
+    _tree(ele)
diff --git a/DrissionPage/_functions/web.pyi b/DrissionPage/_functions/web.pyi
index c592613..aaebf3c 100644
--- a/DrissionPage/_functions/web.pyi
+++ b/DrissionPage/_functions/web.pyi
@@ -11,7 +11,7 @@ from typing import Union
 from requests import Session
 from requests.cookies import RequestsCookieJar
 
-from .._base.base import BasePage, DrissionElement
+from .._base.base import DrissionElement, BaseParser
 from .._elements.chromium_element import ChromiumElement
 from .._pages.chromium_base import ChromiumBase
 
@@ -52,4 +52,4 @@ def is_cookie_in_driver(page: ChromiumBase, cookie: dict) -> bool: ...
 def get_blob(page: ChromiumBase, url: str, as_bytes: bool = True) -> bytes: ...
 
 
-def tree(ele_or_page, layer=5, last_one=False, body='') -> None: ...
+def tree(ele_or_page:BaseParser) -> None: ...
diff --git a/DrissionPage/common.py b/DrissionPage/common.py
index 5c1e2b8..2c09a2e 100644
--- a/DrissionPage/common.py
+++ b/DrissionPage/common.py
@@ -10,7 +10,7 @@ from ._functions.by import By
 from ._functions.keys import Keys
 from ._functions.settings import Settings
 from ._functions.tools import wait_until, configs_to_here
-from ._functions.web import get_blob
+from ._functions.web import get_blob, tree
 from ._units.actions import Actions
 
-__all__ = ['make_session_ele', 'Actions', 'Keys', 'By', 'Settings', 'wait_until', 'configs_to_here', 'get_blob']
+__all__ = ['make_session_ele', 'Actions', 'Keys', 'By', 'Settings', 'wait_until', 'configs_to_here', 'get_blob', 'tree']