feat: 完善类型注释，增加 mypy 类型检测

base/base_crawler.py

@@ -20,10 +20,6 @@ class AbstractLogin(ABC):
     async def begin(self):
         pass
 
-    @abstractmethod
-    async def check_login_state(self):
-        pass
-
     @abstractmethod
     async def login_by_qrcode(self):
         pass

base/proxy_account_pool.py

@@ -1,14 +1,14 @@
-from typing import Tuple, Optional
+from typing import Tuple, Optional, List, Set
 
 import config
 
 
 class PhonePool:
     """phone pool class"""
 
-    def __init__(self):
-        self.phones = []
-        self.used_phones = set()
+    def __init__(self) -> None:
+        self.phones: List[str] = []
+        self.used_phones: Set[str] = set()
 
     def add_phone(self, phone: str) -> bool:
         """add phone to the pool"""
@@ -40,9 +40,9 @@ def clear(self):
 
 
 class IPPool:
-    def __init__(self):
-        self.ips = []
-        self.used_ips = set()
+    def __init__(self) -> None:
+        self.ips: List[str]= []
+        self.used_ips: Set[str]= set()
 
     def add_ip(self, ip):
         """添加ip"""

config/base_config.py

@@ -2,7 +2,7 @@
 PLATFORM = "xhs"
 KEYWORDS = "健身,旅游"
 LOGIN_TYPE = "qrcode"  # qrcode or phone or cookies
-COOKIES = ""
+COOKIES = "web_session=xxxxcfed1566xxxxxxxxxxxxxxxxxxx;" # if platform is xhs, pleas set only web_session cookie attr
 
 # redis config
 REDIS_DB_HOST = "redis://127.0.0.1"  # your redis host
@@ -18,7 +18,7 @@
 HEADLESS = True
 
 # save login state
-SAVE_LOGIN_STATE = True
+SAVE_LOGIN_STATE = False
 
 # save user data dir
 USER_DATA_DIR = "%s_user_data_dir"  # %s will be replaced by platform name

media_platform/douyin/client.py

@@ -18,9 +18,10 @@ def __init__(
             self,
             timeout=30,
             proxies=None,
-            headers: Optional[Dict] = None,
-            playwright_page: Page = None,
-            cookie_dict: Dict = None
+            *,
+            headers: Dict,
+            playwright_page: Optional[Page],
+            cookie_dict: Dict
     ):
         self.proxies = proxies
         self.timeout = timeout
@@ -33,7 +34,7 @@ async def __process_req_params(self, params: Optional[Dict] = None, headers: Opt
         if not params:
             return
         headers = headers or self.headers
-        local_storage: Dict = await self.playwright_page.evaluate("() => window.localStorage")
+        local_storage: Dict = await self.playwright_page.evaluate("() => window.localStorage") # type: ignore
         douyin_js_obj = execjs.compile(open('libs/douyin.js').read())
         common_params = {
             "device_platform": "webapp",
@@ -141,7 +142,7 @@ async def get_video_by_id(self, aweme_id: str):
         del headers["Origin"]
         return await self.get("/aweme/v1/web/aweme/detail/", params, headers)
 
-    async def get_aweme_comments(self, aweme_id: str, cursor: str = ""):
+    async def get_aweme_comments(self, aweme_id: str, cursor: int = 0):
         """get note comments
 
         """

media_platform/douyin/core.py

@@ -20,20 +20,21 @@
 
 
 class DouYinCrawler(AbstractCrawler):
-    def __init__(self):
-        self.browser_context: Optional[BrowserContext] = None
-        self.context_page: Optional[Page] = None
+    dy_client: DOUYINClient
+
+    def __init__(self) -> None:
+        self.browser_context: Optional[BrowserContext] = None # type: ignore
+        self.context_page: Optional[Page] = None # type: ignore
         self.user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/109.0.0.0 Safari/537.36"  # fixed
-        self.dy_client: Optional[DOUYINClient] = None
         self.index_url = "https://www.douyin.com"
-        self.command_args: Optional[Namespace] = None
-        self.account_pool: Optional[AccountPool] = None
+        self.command_args: Optional[Namespace] = None # type: ignore
+        self.account_pool: Optional[AccountPool] = None # type: ignore
 
     def init_config(self, **kwargs):
         for key, value in kwargs.items():
             setattr(self, key, value)
 
-    async def start(self):
+    async def start(self) -> None:
         account_phone, playwright_proxy, httpx_proxy = self.create_proxy_info()
         async with async_playwright() as playwright:
             # Launch a browser context.
@@ -52,7 +53,7 @@ async def start(self):
             self.dy_client = await self.create_douyin_client(httpx_proxy)
             if not await self.dy_client.ping(browser_context=self.browser_context):
                 login_obj = DouYinLogin(
-                    login_type=self.command_args.lt,
+                    login_type=self.command_args.lt, # type: ignore
                     login_phone=account_phone,
                     browser_context=self.browser_context,
                     context_page=self.context_page,
@@ -66,7 +67,7 @@ async def start(self):
 
             utils.logger.info("Douyin Crawler finished ...")
 
-    async def search_posts(self):
+    async def search_posts(self) -> None:
         utils.logger.info("Begin search douyin keywords")
         for keyword in config.KEYWORDS.split(","):
             utils.logger.info(f"Current keyword: {keyword}")
@@ -87,7 +88,7 @@ async def search_posts(self):
                                            post_item.get("aweme_mix_info", {}).get("mix_items")[0]
                     except TypeError:
                         continue
-                    aweme_list.append(aweme_info.get("aweme_id"))
+                    aweme_list.append(aweme_info.get("aweme_id",""))
                     await douyin.update_douyin_aweme(aweme_item=aweme_info)
             utils.logger.info(f"keyword:{keyword}, aweme_list:{aweme_list}")
             # await self.batch_get_note_comments(aweme_list)
@@ -115,7 +116,7 @@ def create_proxy_info(self) -> Tuple[Optional[str], Optional[Dict], Optional[str
             return None, None, None
 
         # phone: 13012345671  ip_proxy: 111.122.xx.xx1:8888
-        phone, ip_proxy = self.account_pool.get_account()
+        phone, ip_proxy = self.account_pool.get_account() # type: ignore
         playwright_proxy = {
             "server": f"{config.IP_PROXY_PROTOCOL}{ip_proxy}",
             "username": config.IP_PROXY_USER,
@@ -124,9 +125,9 @@ def create_proxy_info(self) -> Tuple[Optional[str], Optional[Dict], Optional[str
         httpx_proxy = f"{config.IP_PROXY_PROTOCOL}{config.IP_PROXY_USER}:{config.IP_PROXY_PASSWORD}@{ip_proxy}"
         return phone, playwright_proxy, httpx_proxy
 
-    async def create_douyin_client(self, httpx_proxy: str) -> DOUYINClient:
+    async def create_douyin_client(self, httpx_proxy: Optional[str]) -> DOUYINClient:
         """Create douyin client"""
-        cookie_str, cookie_dict = utils.convert_cookies(await self.browser_context.cookies())
+        cookie_str, cookie_dict = utils.convert_cookies(await self.browser_context.cookies()) # type: ignore
         douyin_client = DOUYINClient(
             proxies=httpx_proxy,
             headers={
@@ -151,18 +152,18 @@ async def launch_browser(
     ) -> BrowserContext:
         """Launch browser and create browser context"""
         if config.SAVE_LOGIN_STATE:
-            user_data_dir = os.path.join(os.getcwd(), "browser_data", config.USER_DATA_DIR % self.command_args.platform)
+            user_data_dir = os.path.join(os.getcwd(), "browser_data", config.USER_DATA_DIR % self.command_args.platform) # type: ignore
             browser_context = await chromium.launch_persistent_context(
                 user_data_dir=user_data_dir,
                 accept_downloads=True,
                 headless=headless,
-                proxy=playwright_proxy,
+                proxy=playwright_proxy, # type: ignore
                 viewport={"width": 1920, "height": 1080},
                 user_agent=user_agent
-            )
+            ) # type: ignore
             return browser_context
         else:
-            browser = await chromium.launch(headless=headless, proxy=playwright_proxy)
+            browser = await chromium.launch(headless=headless, proxy=playwright_proxy) # type: ignore
             browser_context = await browser.new_context(
                 viewport={"width": 1920, "height": 1080},
                 user_agent=user_agent

media_platform/douyin/login.py

@@ -1,6 +1,7 @@
 import sys
 import asyncio
 import functools
+from typing import Optional
 
 import aioredis
 from tenacity import (
@@ -22,10 +23,10 @@ class DouYinLogin(AbstractLogin):
 
     def __init__(self,
                  login_type: str,
-                 browser_context: BrowserContext,
-                 context_page: Page,
-                 login_phone: str = None,
-                 cookie_str: str = None
+                 browser_context: BrowserContext, # type: ignore
+                 context_page: Page, # type: ignore
+                 login_phone: Optional[str] = "",
+                 cookie_str: Optional[str] = ""
                  ):
         self.login_type = login_type
         self.browser_context = browser_context
@@ -202,14 +203,14 @@ async def move_slider(self, back_selector: str, gap_selector: str, move_step: in
             selector=back_selector,
             timeout=1000 * 10,  # wait 10 seconds
         )
-        slide_back = str(await slider_back_elements.get_property("src"))
+        slide_back = str(await slider_back_elements.get_property("src")) # type: ignore
 
         # get slider gap image
         gap_elements = await self.context_page.wait_for_selector(
             selector=gap_selector,
             timeout=1000 * 10,  # wait 10 seconds
         )
-        gap_src = str(await gap_elements.get_property("src"))
+        gap_src = str(await gap_elements.get_property("src")) # type: ignore
 
         # 识别滑块位置
         slide_app = utils.Slide(gap=gap_src, bg=slide_back)
@@ -223,14 +224,14 @@ async def move_slider(self, back_selector: str, gap_selector: str, move_step: in
 
         # 根据轨迹拖拽滑块到指定位置
         element = await self.context_page.query_selector(gap_selector)
-        bounding_box = await element.bounding_box()
+        bounding_box = await element.bounding_box() # type: ignore
 
-        await self.context_page.mouse.move(bounding_box["x"] + bounding_box["width"] / 2,
-                                           bounding_box["y"] + bounding_box["height"] / 2)
+        await self.context_page.mouse.move(bounding_box["x"] + bounding_box["width"] / 2, # type: ignore
+                                           bounding_box["y"] + bounding_box["height"] / 2) # type: ignore
         # 这里获取到x坐标中心点位置
-        x = bounding_box["x"] + bounding_box["width"] / 2
+        x = bounding_box["x"] + bounding_box["width"] / 2 # type: ignore
         # 模拟滑动操作
-        await element.hover()
+        await element.hover() # type: ignore
         await self.context_page.mouse.down()
 
         for track in tracks:

media_platform/xhs/client.py

@@ -17,9 +17,10 @@ def __init__(
             self,
             timeout=10,
             proxies=None,
-            headers: Optional[Dict] = None,
-            playwright_page: Page = None,
-            cookie_dict: Dict = None
+            *,
+            headers: Dict[str, str],
+            playwright_page: Page,
+            cookie_dict: Dict[str, str],
     ):
         self.proxies = proxies
         self.timeout = timeout
@@ -51,29 +52,29 @@ async def _pre_headers(self, url: str, data=None):
         self.headers.update(headers)
         return self.headers
 
-    async def request(self, method, url, **kwargs):
+    async def request(self, method, url, **kwargs) -> Dict:
         async with httpx.AsyncClient(proxies=self.proxies) as client:
             response = await client.request(
                 method, url, timeout=self.timeout,
                 **kwargs
             )
-        data = response.json()
+        data: Dict = response.json()
         if data["success"]:
-            return data.get("data", data.get("success"))
+            return data.get("data", data.get("success", {}))
         elif data["code"] == self.IP_ERROR_CODE:
             raise IPBlockError(self.IP_ERROR_STR)
         else:
             raise DataFetchError(data.get("msg", None))
 
-    async def get(self, uri: str, params=None):
+    async def get(self, uri: str, params=None) -> Dict:
         final_uri = uri
         if isinstance(params, dict):
             final_uri = (f"{uri}?"
                          f"{'&'.join([f'{k}={v}' for k, v in params.items()])}")
         headers = await self._pre_headers(final_uri)
         return await self.request(method="GET", url=f"{self._host}{final_uri}", headers=headers)
 
-    async def post(self, uri: str, data: dict):
+    async def post(self, uri: str, data: dict) -> Dict:
         headers = await self._pre_headers(uri, data)
         json_str = json.dumps(data, separators=(',', ':'), ensure_ascii=False)
         return await self.request(method="POST", url=f"{self._host}{uri}",
@@ -86,7 +87,7 @@ async def ping(self) -> bool:
         try:
             note_card: Dict = await self.get_note_by_id(note_id)
             return note_card.get("note_id") == note_id
-        except DataFetchError:
+        except Exception:
             return False
 
     async def update_cookies(self, browser_context: BrowserContext):
@@ -128,7 +129,8 @@ async def get_note_by_id(self, note_id: str) -> Dict:
         data = {"source_note_id": note_id}
         uri = "/api/sns/web/v1/feed"
         res = await self.post(uri, data)
-        return res["items"][0]["note_card"]
+        res_dict: Dict = res["items"][0]["note_card"]
+        return res_dict
 
     async def get_note_comments(self, note_id: str, cursor: str = "") -> Dict:
         """get note comments

media_platform/xhs/core.py

@@ -21,15 +21,15 @@
 
 
 class XiaoHongShuCrawler(AbstractCrawler):
+    context_page: Page
+    browser_context: BrowserContext
+    xhs_client: XHSClient
+    account_pool: AccountPool
 
     def __init__(self):
-        self.browser_context: Optional[BrowserContext] = None
-        self.context_page: Optional[Page] = None
-        self.user_agent = utils.get_user_agent()
-        self.xhs_client: Optional[XHSClient] = None
         self.index_url = "https://www.xiaohongshu.com"
-        self.command_args: Optional[Namespace] = None
-        self.account_pool: Optional[AccountPool] = None
+        self.command_args: Optional[Namespace] = None # type: ignore
+        self.user_agent = utils.get_user_agent()
 
     def init_config(self, **kwargs):
         for key, value in kwargs.items():
@@ -69,7 +69,7 @@ async def start(self):
 
             utils.logger.info("Xhs Crawler finished ...")
 
-    async def search_posts(self):
+    async def search_posts(self) -> None:
         """Search for notes and retrieve their comment information."""
         utils.logger.info("Begin search xiaohongshu keywords")
 
@@ -86,7 +86,7 @@ async def search_posts(self):
                 _semaphore = asyncio.Semaphore(config.MAX_CONCURRENCY_NUM)
                 task_list = [
                     self.get_note_detail(post_item.get("id"), _semaphore)
-                    for post_item in posts_res.get("items")
+                    for post_item in posts_res.get("items", {})
                 ]
                 note_details = await asyncio.gather(*task_list)
                 for note_detail in note_details:
@@ -170,18 +170,18 @@ async def launch_browser(
         if config.SAVE_LOGIN_STATE:
             # feat issue #14
             # we will save login state to avoid login every time
-            user_data_dir = os.path.join(os.getcwd(), "browser_data", config.USER_DATA_DIR % self.command_args.platform)
+            user_data_dir = os.path.join(os.getcwd(), "browser_data", config.USER_DATA_DIR % self.command_args.platform) # type: ignore
             browser_context = await chromium.launch_persistent_context(
                 user_data_dir=user_data_dir,
                 accept_downloads=True,
                 headless=headless,
-                proxy=playwright_proxy,
+                proxy=playwright_proxy, # type: ignore
                 viewport={"width": 1920, "height": 1080},
                 user_agent=user_agent
             )
             return browser_context
         else:
-            browser = await chromium.launch(headless=headless, proxy=playwright_proxy)
+            browser = await chromium.launch(headless=headless, proxy=playwright_proxy) # type: ignore
             browser_context = await browser.new_context(
                 viewport={"width": 1920, "height": 1080},
                 user_agent=user_agent

media_platform/xhs/help.py

@@ -92,7 +92,7 @@ def mrc(e):
     ]
     o = -1
 
-    def right_without_sign(num, bit=0) -> int:
+    def right_without_sign(num: int, bit: int=0) -> int:
         val = ctypes.c_uint32(num).value >> bit
         MAX32INT = 4294967295
         return (val + (MAX32INT + 1)) % (2 * (MAX32INT + 1)) - MAX32INT - 1