From e5e7c57b32558c5e004a375ace248c7ac9187cff Mon Sep 17 00:00:00 2001 From: g1879 Date: Thu, 3 Dec 2020 17:13:02 +0800 Subject: [PATCH] =?UTF-8?q?=E5=9F=BA=E6=9C=AC=E5=AE=8C=E6=88=90cookies?= =?UTF-8?q?=E7=9B=B8=E5=85=B3=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- DrissionPage/config.py | 57 ++++++++++++++++++++++++++++++++-------- DrissionPage/drission.py | 37 ++++++++++---------------- 2 files changed, 60 insertions(+), 34 deletions(-) diff --git a/DrissionPage/config.py b/DrissionPage/config.py index 38c5cf3..3d331c1 100644 --- a/DrissionPage/config.py +++ b/DrissionPage/config.py @@ -358,8 +358,8 @@ class SessionOptions(object): """ self._max_redirects = max_redirects - def set_header(self, attr: str, value: str): - """设置header中一个项 \n + def set_a_header(self, attr: str, value: str): + """设置headers中一个项 \n :param attr: 设置名称 :param value: 设置值 :return: 返回当前对象 @@ -370,7 +370,7 @@ class SessionOptions(object): self._headers[attr.lower()] = value return self - def remove_header(self, attr: str): + def remove_a_header(self, attr: str): """从headers中删除一个设置 \n :param attr: 要删除的设置 :return: 返回当前对象 @@ -700,13 +700,10 @@ def _session_options_to_dict(options: Union[dict, SessionOptions, None]) -> Unio re_dict = dict() attrs = ['headers', 'proxies', 'hooks', 'params', 'verify', 'stream', 'trust_env', 'max_redirects'] # 'adapters', - val = options.__getattribute__('_cookies') + cookies = options.__getattribute__('_cookies') - if val is not None: - if isinstance(val, RequestsCookieJar): - re_dict['cookies'] = [_cookie_to_dict(cookie) for cookie in val] - else: - re_dict['cookies'] = val + if cookies is not None: + re_dict['cookies'] = _cookies_to_tuple(cookies) for attr in attrs: val = options.__getattribute__(f'_{attr}') @@ -720,7 +717,7 @@ def _session_options_to_dict(options: Union[dict, SessionOptions, None]) -> Unio return re_dict -def _cookie_to_dict(cookie: Cookie) -> dict: +def _cookie_to_dict(cookie: Union[Cookie, str, dict]) -> dict: """把Cookie对象转为dict格式 \n :param cookie: Cookie对象 :return: cookie字典 @@ -732,6 +729,44 @@ def _cookie_to_dict(cookie: Cookie) -> dict: return cookie_dict elif isinstance(cookie, dict): - return cookie + cookie_dict = cookie + + elif isinstance(cookie, str): + cookie = cookie.split(';') + cookie_dict = {} + + for key, attr in enumerate(cookie): + attr_val = attr.lstrip().split('=') + + if key == 0: + cookie_dict['name'] = attr_val[0] + cookie_dict['value'] = attr_val[1] + else: + cookie_dict[attr_val[0]] = attr_val[1] + + return cookie_dict + else: raise TypeError + + return cookie_dict + + +def _cookies_to_tuple(cookies: Union[RequestsCookieJar, list, tuple, str, dict]) -> tuple: + """把cookies转为tuple格式 \n + :param cookies: cookies信息，可为CookieJar, list, tuple, str, dict + :return: 返回tuple形式的cookies + """ + if isinstance(cookies, (list, tuple, RequestsCookieJar)): + cookies = tuple(_cookie_to_dict(cookie) for cookie in cookies) + + elif isinstance(cookies, str): + cookies = tuple(dict([cookie.lstrip().split("=", 1)]) for cookie in cookies.split(";")) + + elif isinstance(cookies, dict): + cookies = tuple({'name': cookie, 'value': cookies[cookie]} for cookie in cookies) + + else: + raise TypeError + + return cookies diff --git a/DrissionPage/drission.py b/DrissionPage/drission.py index 2335225..0b88c50 100644 --- a/DrissionPage/drission.py +++ b/DrissionPage/drission.py @@ -10,13 +10,13 @@ from typing import Union from requests import Session from requests.cookies import RequestsCookieJar from selenium import webdriver -from selenium.common.exceptions import SessionNotCreatedException, UnableToSetCookieException +from selenium.common.exceptions import SessionNotCreatedException from selenium.webdriver.chrome.options import Options from selenium.webdriver.chrome.webdriver import WebDriver from tldextract import extract from .config import (_dict_to_chrome_options, _session_options_to_dict, - SessionOptions, DriverOptions, _chrome_options_to_dict, OptionsManager, _cookie_to_dict) + SessionOptions, DriverOptions, _chrome_options_to_dict, OptionsManager, _cookies_to_tuple) class Drission(object): @@ -164,19 +164,13 @@ class Drission(object): :param set_driver: 是否设置driver的cookies :return: None """ - if isinstance(cookies, (list, tuple, RequestsCookieJar)): - cookies = tuple(_cookie_to_dict(cookie) for cookie in cookies) - elif isinstance(cookies, str): - cookies = tuple(dict([cookie.lstrip().split("=", 1)]) for cookie in cookies.split(";")) - elif isinstance(cookies, dict): - cookies = tuple({'name': cookie, 'value': cookies[cookie]} for cookie in cookies) - else: - raise TypeError + cookies = _cookies_to_tuple(cookies) for cookie in cookies: if cookie['value'] is None: cookie['value'] = '' + # 添加cookie到session if set_session: kwargs = {x: cookie[x] for x in cookie if x not in ('name', 'value', 'httpOnly', 'expiry')} @@ -185,26 +179,23 @@ class Drission(object): self.session.cookies.set(cookie['name'], cookie['value'], **kwargs) + # 添加cookie到driver if set_driver: if 'expiry' in cookie: cookie['expiry'] = int(cookie['expiry']) + cookie_domain = cookie['domain'] if cookie['domain'][0] != '.' else cookie['domain'][1:] + try: - self.driver.add_cookie(cookie) + browser_domain = extract(self.driver.current_url).fqdn + except AttributeError: + browser_domain = '' - except UnableToSetCookieException: - cookie_domain = cookie['domain'] if cookie['domain'][0] != '.' else cookie['domain'][1:] + if cookie_domain not in browser_domain: + self.driver.get(cookie_domain if cookie_domain.startswith('http://') + else f'http://{cookie_domain}') - try: - browser_domain = extract(self.driver.current_url).fqdn - except AttributeError: - browser_domain = '' - - if cookie_domain not in browser_domain: - self.driver.get(cookie_domain if cookie_domain.startswith('http://') - else f'http://{cookie_domain}') - - self.driver.add_cookie(cookie) + self.driver.add_cookie(cookie) def _set_session(self, data: dict) -> None: if self._session is None: