基本完善cookies方面功能

DrissionPage/drission.py

@@ -10,7 +10,7 @@ from typing import Union
 from requests import Session
 from requests.cookies import RequestsCookieJar
 from selenium import webdriver
-from selenium.common.exceptions import SessionNotCreatedException, UnableToSetCookieException  # , WebDriverException
+from selenium.common.exceptions import SessionNotCreatedException, UnableToSetCookieException
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.chrome.webdriver import WebDriver
 from tldextract import extract
@@ -19,9 +19,6 @@ from .config import (_dict_to_chrome_options, _session_options_to_dict,
                      SessionOptions, DriverOptions, _chrome_options_to_dict, OptionsManager, _cookie_to_dict)
 
 
-# from urllib.parse import urlparse
-
-
 class Drission(object):
     """Drission类用于管理WebDriver对象和Session对象，是驱动器的角色"""
 
@@ -155,18 +152,17 @@ class Drission(object):
             self._driver.get(url)
 
             for cookie in cookies:
-                # self._ensure_add_cookie(cookie)
                 self.set_cookies(cookie, set_driver=True)
 
     def set_cookies(self,
-                    cookies: Union[RequestsCookieJar, list, tuple, str],
+                    cookies: Union[RequestsCookieJar, list, tuple, str, dict],
                     set_session: bool = False,
                     set_driver: bool = False) -> None:
-        """
+        """设置cookies                                                      \n
-        :param cookies:
+        :param cookies: cookies信息，可为CookieJar, list, tuple, str, dict
-        :param set_session:
+        :param set_session: 是否设置session的cookies
-        :param set_driver:
+        :param set_driver: 是否设置driver的cookies
-        :return:
+        :return: None
         """
         if isinstance(cookies, (list, tuple, RequestsCookieJar)):
             cookies = tuple(_cookie_to_dict(cookie) for cookie in cookies)
@@ -190,6 +186,8 @@ class Drission(object):
                 self.session.cookies.set(cookie['name'], cookie['value'], **kwargs)
 
             if set_driver:
+                if 'expiry' in cookie:
+                    cookie['expiry'] = int(cookie['expiry'])
 
                 try:
                     self.driver.add_cookie(cookie)
@@ -203,25 +201,11 @@ class Drission(object):
                         browser_domain = ''
 
                     if cookie_domain not in browser_domain:
-                        self.driver.get(f'http://{cookie_domain.lstrip("http://")}')
+                        self.driver.get(cookie_domain if cookie_domain.startswith('http://')
+                                        else f'http://{cookie_domain}')
 
                     self.driver.add_cookie(cookie)
 
-    def add_a_cookie(self,
-                     cookie: str,
-                     set_session: bool = False,
-                     set_driver: bool = False):
-        pass
-
-    def remove_a_cookie(self,
-                        name: str,
-                        set_session: bool = False,
-                        set_driver: bool = False):
-        pass
-
-    def clear_cookies(self):
-        pass
-
     def _set_session(self, data: dict) -> None:
         if self._session is None:
             self._session = Session()
@@ -245,8 +229,6 @@ class Drission(object):
             self.user_agent_to_session(self.driver, self.session)
 
         self.set_cookies(self.driver.get_cookies(), set_session=True)
-        # for cookie in driver.get_cookies():
-        # session.cookies.set(cookie['name'], cookie['value'], domain=cookie['domain'])
 
     def cookies_to_driver(self, url: str) -> None:
         """把session对象的cookies复制到driver对象  \n
@@ -259,80 +241,6 @@ class Drission(object):
 
         self.set_cookies(cookies, set_driver=True)
 
-    # def cookies_to_driver(self, url: str,
-    #                       driver: WebDriver = None,
-    #                       session: Session = None) -> None:
-    #     """把session对象的cookies复制到driver对象  \n
-    #     :param url: 作用域
-    #     :param driver: 目标driver对象
-    #     :param session: 来源session对象
-    #     :return: None
-    #     """
-    #     driver = driver or self.driver
-    #     session = session or self.session
-    #     domain = urlparse(url).netloc
-    #
-    #     if not domain:
-    #         raise Exception('Without specifying a domain')
-    #
-    #     # 翻译cookies
-    #     for i in [x for x in session.cookies if domain in x.domain]:
-    #         cookie_data = _cookie_to_dict(i)
-    #         self._ensure_add_cookie(cookie_data, driver=driver)
-    #
-    # def _ensure_add_cookie(self, cookie, override_domain=None, driver=None) -> None:
-    #     """添加cookie到driver                  \n
-    #     :param cookie: 要添加的cookie
-    #     :param override_domain: 覆盖作用域
-    #     :param driver: 操作的driver对象
-    #     :return: None
-    #     """
-    #     driver = driver or self.driver
-    #
-    #     if override_domain:
-    #         cookie['domain'] = override_domain
-    #
-    #     cookie_domain = cookie['domain'] if cookie['domain'][0] != '.' else cookie['domain'][1:]
-    #
-    #     try:
-    #         browser_domain = extract(driver.current_url).fqdn
-    #     except AttributeError:
-    #         browser_domain = ''
-    #
-    #     if cookie_domain not in browser_domain:
-    #         driver.get(f'http://{cookie_domain.lstrip("http://")}')
-    #
-    #     if 'expiry' in cookie:
-    #         cookie['expiry'] = int(cookie['expiry'])
-    #
-    #     driver.add_cookie(cookie)
-    #
-    #     # 如果添加失败，尝试更宽的域名
-    #     if not self._is_cookie_in_driver(cookie, driver):
-    #         cookie['domain'] = extract(cookie['domain']).registered_domain
-    #         driver.add_cookie(cookie)
-    #
-    #         if not self._is_cookie_in_driver(cookie):
-    #             raise WebDriverException(f"Couldn't add the following cookie to the webdriver\n{cookie}\n")
-    #
-    # def _is_cookie_in_driver(self, cookie, driver=None) -> bool:
-    #     """检查cookie是否已经在driver里                   \n
-    #     只检查name、value、domain，检查domain时比较宽      \n
-    #     :param cookie: 要检查的cookie
-    #     :param driver: 被检查的driver
-    #     :return: 返回布尔值
-    #     """
-    #     driver = driver or self.driver
-    #     for driver_cookie in driver.get_cookies():
-    #
-    #         if (cookie['name'] == driver_cookie['name'] and
-    #                 cookie['value'] == driver_cookie['value'] and
-    #                 (cookie['domain'] == driver_cookie['domain'] or
-    #                  f'.{cookie["domain"]}' == driver_cookie['domain'])):
-    #             return True
-    #
-    #     return False
-
     def user_agent_to_session(self, driver: WebDriver = None, session: Session = None) -> None:
         """把driver的user-agent复制到session    \n
         :param driver: 来源driver对象

DrissionPage/mix_page.py

@@ -7,6 +7,7 @@
 from typing import Union, List, Tuple
 
 from requests import Response, Session
+from requests.cookies import RequestsCookieJar
 from selenium.webdriver.chrome.webdriver import WebDriver
 from selenium.webdriver.remote.webelement import WebElement
 
@@ -139,10 +140,24 @@ class MixPage(Null, SessionPage, DriverPage):
         elif self._mode == 'd':
             return super(SessionPage, self).title
 
-    def get_cookies(self, as_dict: bool = False) -> Union[dict, list]:
+    def set_cookies(self, cookies: Union[RequestsCookieJar, list, tuple, str, dict]) -> None:
-        """返回cookies"""
+        """设置cookies                                                          \n
+        :param cookies: cookies信息，可为CookieJar, list, tuple, str, dict
+        :return: None
+        """
         if self._mode == 's':
-            return super().get_cookies(as_dict)
+            self.drission.set_cookies(cookies, set_session=True)
+        elif self._mode == 'd':
+            self.drission.set_cookies(cookies, set_driver=True)
+
+    def get_cookies(self, as_dict: bool = False, all_domains: bool = False) -> Union[dict, list]:
+        """返回cookies                               \n
+        :param as_dict: 是否以字典方式返回
+        :param all_domains: 是否返回所有域的cookies
+        :return: cookies信息
+        """
+        if self._mode == 's':
+            return super().get_cookies(as_dict, all_domains)
         elif self._mode == 'd':
             return super(SessionPage, self).get_cookies(as_dict)
 
@@ -162,8 +177,10 @@ class MixPage(Null, SessionPage, DriverPage):
         if self._mode == 'd':
             self._driver = True
             self._url = None if not self._driver else self._drission.driver.current_url
+
             if self._session_url:
                 self.cookies_to_driver(self._session_url)
+
                 if go:
                     self.get(self._session_url)
 
@@ -171,8 +188,10 @@ class MixPage(Null, SessionPage, DriverPage):
         elif self._mode == 's':
             self._session = True
             self._url = self._session_url
+
             if self._driver:
                 self.cookies_to_session()
+
                 if go and self._drission.driver.current_url.startswith('http'):
                     self.get(self._drission.driver.current_url)
 

DrissionPage/session_page.py

@@ -15,6 +15,7 @@ from typing import Union, List, Tuple
 from urllib.parse import urlparse, quote, unquote
 
 from requests import Session, Response
+from tldextract import extract
 
 from .common import str_to_loc, translate_loc, get_available_file_name, format_html
 from .config import _cookie_to_dict
@@ -67,12 +68,23 @@ class SessionPage(object):
         """返回页面html文本"""
         return format_html(self.response.text)
 
-    def get_cookies(self, as_dict: bool = False) -> Union[dict, list]:
+    def get_cookies(self, as_dict: bool = False, all_domains: bool = False) -> Union[dict, list]:
-        """返回session的cookies"""
+        """返回cookies                               \n
-        if as_dict:
+        :param as_dict: 是否以字典方式返回
-            return self.session.cookies.get_dict()
+        :param all_domains: 是否返回所有域的cookies
+        :return: cookies信息
+        """
+        if all_domains:
+            cookies = self.session.cookies
         else:
-            return [_cookie_to_dict(cookie) for cookie in self.session.cookies]
+            url = extract(self.url)
+            domain = f'{url.domain}.{url.suffix}'
+            cookies = tuple(x for x in self.session.cookies if domain in x.domain)
+
+        if as_dict:
+            return {x.name: x.value for x in cookies}
+        else:
+            return [_cookie_to_dict(cookie) for cookie in cookies]
 
     def ele(self,
             loc_or_ele: Union[Tuple[str, str], str, SessionElement],

README.en.md

@@ -547,6 +547,7 @@ page.current_tab_handle  # Return to the current tab page handle
 When calling a method that only belongs to d mode, it will automatically switch to d mode. See APIs for detailed usage.
 
 ```python
+page.set_cookies()  # set cookies
 page.get_cookies() # Get cookies, which can be returned by list or dict
 page.change_mode()  # Switch mode, it will automatically copy cookies
 page.cookies_to_session()  # Copy cookies from WebDriver object to Session object
@@ -1196,6 +1197,22 @@ Returns: None
 
 
 
+### set_cookies()
+
+Set cookies.
+
+Parameter Description:
+
+- cookies: Union[RequestsCookieJar, list, tuple, str, dict]-cookies information, can be CookieJar, list, tuple, str, dict
+
+- set_session: bool-whether to set session cookies
+
+- set_driver: bool-whether to set driver cookies
+
+Returns: None
+
+
+
 ### cookies_to_driver()
 
 Copy cookies from session to driver.
@@ -1341,6 +1358,18 @@ Returns: bool
 
 
 
+### set_cookies()
+
+Set cookies.
+
+Parameter Description:
+
+- cookies: Union[RequestsCookieJar, list, tuple, str, dict]  - cookies information, can be CookieJar, list, tuple, str, dict
+
+Returns: None
+
+
+
 ### get_cookies()
 
 Return cookies.
@@ -1349,6 +1378,10 @@ Parameter Description:
 
 - as_dict: bool  - Whether to return as dict, the default is to return complete cookies as list
 
+- all_domains: bool  - whether to return cookies of all domains, only valid in s mode
+
+Returns: a dictionary or list of cookies
+
 
 
 ### change_mode()

README.zh-cn.md

@@ -551,6 +551,7 @@ page.current_tab_handle  # 返回当前标签页 handle
 调用只属于 d 模式的方法，会自动切换到 d 模式。详细用法见 APIs。
 
 ```python
+page.set_cookies()  # 设置cookies
 page.get_cookies()  # 获取 cookies，可以 list 或 dict 方式返回
 page.change_mode()  # 切换模式，会自动复制 cookies
 page.cookies_to_session()  # 从 WebDriver 对象复制 cookies 到 Session 对象
@@ -1201,6 +1202,20 @@ Drission 类用于管理 WebDriver 对象和 Session 对象，是驱动器的角
 
 
 
+### set_cookies()
+
+设置 cookies。
+
+参数说明：
+
+- cookies: Union[RequestsCookieJar, list, tuple, str, dict]  - cookies 信息，可为CookieJar, list, tuple, str, dict
+- set_session: bool                                                             - 是否设置 session 的 cookies
+- set_driver: bool                                                                - 是否设置 driver 的 cookies
+
+返回： None
+
+
+
 ### cookies_to_session()
 
 把 driver 对象的 cookies 复制到 session 对象。
@@ -1358,6 +1373,18 @@ MixPage 封装了页面操作的常用功能，可在 driver 和 session 模式
 
 
 
+### set_cookies()
+
+设置 cookies。
+
+参数说明：
+
+- cookies: Union[RequestsCookieJar, list, tuple, str, dict]  - cookies 信息，可为CookieJar, list, tuple, str, dict
+
+返回： None
+
+
+
 ### get_cookies()
 
 返回 cookies。
@@ -1365,6 +1392,9 @@ MixPage 封装了页面操作的常用功能，可在 driver 和 session 模式
 参数说明：
 
 - as_dict: bool          - 是否以 dict 方式返回，默认以 list 返回完整的 cookies
+- all_domains: bool  - 是否返回所有域名的 cookies，只有 s 模式下生效
+
+返回：cookies 字典或列表