From b2da0a902d8b8f69ef84a8d03180fec132e72867 Mon Sep 17 00:00:00 2001
From: HibiKier <775757368@qq.com>
Date: Thu, 7 Nov 2024 13:38:20 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=A8=20playwright=E6=B7=BB=E5=8A=A0cook?=
 =?UTF-8?q?ie=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md                      |  2 +-
 zhenxun/models/chat_history.py | 12 +++++-------
 zhenxun/utils/http_utils.py    | 14 ++++++++++++--
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 83ec4ea8..d7c24a68 100644
--- a/README.md
+++ b/README.md
@@ -60,7 +60,7 @@
 
 <img width="350" height="350" src="https://raw.githubusercontent.com/HibiKier/zhenxun_bot/main/docs_image/help.png"/>
 <img width="250" height="500" src="https://raw.githubusercontent.com/HibiKier/zhenxun_bot/main/docs_image/html_help.png"/>
-<img width="180" height="450" src="https://github.com/HibiKier/zhenxun_bot/blob/dev/docs_image/zhenxun_help.png"/>
+<img width="180" height="450" src="https://raw.githubusercontent.com/HibiKier/zhenxun_bot/main/docs_image/zhenxun_help.png"/>
 
 </div>
 
diff --git a/zhenxun/models/chat_history.py b/zhenxun/models/chat_history.py
index 02425987..dd603683 100644
--- a/zhenxun/models/chat_history.py
+++ b/zhenxun/models/chat_history.py
@@ -1,5 +1,5 @@
 from datetime import datetime, timedelta
-from typing import Literal, Tuple
+from typing import Literal
 
 from tortoise import fields
 from tortoise.functions import Count
@@ -27,7 +27,7 @@ class ChatHistory(Model):
     platform = fields.CharField(255, null=True)
     """平台"""
 
-    class Meta:
+    class Meta: # type: ignore
         table = "chat_history"
         table_description = "聊天记录数据表"
 
@@ -53,7 +53,7 @@ class ChatHistory(Model):
             query = query.filter(create_time__range=date_scope)
         return list(
             await query.annotate(count=Count("user_id"))
-            .order_by(o + "count")
+            .order_by(f"{o}count")
             .group_by("user_id")
             .limit(limit)
             .values_list("user_id", "count")
@@ -74,9 +74,7 @@ class ChatHistory(Model):
             )
         else:
             message = await cls.all().order_by("create_time").first()
-        if message:
-            return message.create_time
-        return None
+        return message.create_time if message else None
 
     @classmethod
     async def get_message(
@@ -85,7 +83,7 @@ class ChatHistory(Model):
         gid: str,
         type_: Literal["user", "group"],
         msg_type: Literal["private", "group"] | None = None,
-        days: int | Tuple[datetime, datetime] | None = None,
+        days: int | tuple[datetime, datetime] | None = None,
     ) -> list[Self]:
         """获取消息查询query
 
diff --git a/zhenxun/utils/http_utils.py b/zhenxun/utils/http_utils.py
index c0e6a74b..dc62f760 100644
--- a/zhenxun/utils/http_utils.py
+++ b/zhenxun/utils/http_utils.py
@@ -431,14 +431,20 @@ class AsyncHttpx:
 class AsyncPlaywright:
     @classmethod
     @asynccontextmanager
-    async def new_page(cls, **kwargs) -> AsyncGenerator[Page, None]:
+    async def new_page(
+        cls, cookies: list[dict[str, Any]] | dict[str, Any] | None = None, **kwargs
+    ) -> AsyncGenerator[Page, None]:
         """获取一个新页面
 
         参数:
-            user_agent: 请求头
+            cookies: cookies
         """
         browser = await get_browser()
         ctx = await browser.new_context(**kwargs)
+        if cookies:
+            if isinstance(cookies, dict):
+                cookies = [cookies]
+            await ctx.add_cookies(cookies)  # type: ignore
         page = await ctx.new_page()
         try:
             yield page
@@ -461,6 +467,7 @@ class AsyncPlaywright:
         timeout: float | None = None,
         type_: Literal["jpeg", "png"] | None = None,
         user_agent: str | None = None,
+        cookies: list[dict[str, Any]] | dict[str, Any] | None = None,
         **kwargs,
     ) -> UniMessage | None:
         """截图，该方法仅用于简单快捷截图，复杂截图请操作 page
@@ -474,6 +481,8 @@ class AsyncPlaywright:
             wait_until: 等待类型
             timeout: 超时限制
             type_: 保存类型
+            user_agent: user_agent
+            cookies: cookies
         """
         if viewport_size is None:
             viewport_size = {"width": 2560, "height": 1080}
@@ -482,6 +491,7 @@ class AsyncPlaywright:
         wait_time = wait_time * 1000 if wait_time else None
         element_list = [element] if isinstance(element, str) else element
         async with cls.new_page(
+            cookies,
             viewport=viewport_size,
             user_agent=user_agent,
             **kwargs,