From b04cbc2ced173fe90b20bf1eab7eabefd1006b1f Mon Sep 17 00:00:00 2001
From: suyiiyii <suyiiyii@gmail.com>
Date: Mon, 23 Sep 2024 13:49:39 +0800
Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E5=AF=BC=E5=85=A5=E5=AF=BC?=
 =?UTF-8?q?=E5=87=BA=20cookie=20=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../config/subs_io/nbesf_model/v3.py          |  9 +--------
 nonebot_bison/config/subs_io/subs_io.py       | 19 ++++++-------------
 nonebot_bison/config/subs_io/utils.py         |  8 ++++++++
 3 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/nonebot_bison/config/subs_io/nbesf_model/v3.py b/nonebot_bison/config/subs_io/nbesf_model/v3.py
index 846c4ff..a2d5e42 100644
--- a/nonebot_bison/config/subs_io/nbesf_model/v3.py
+++ b/nonebot_bison/config/subs_io/nbesf_model/v3.py
@@ -111,14 +111,7 @@ async def magic_cookie_gen(nbesf_data: SubGroup):
     logger.info("开始添加 Cookie 流程")
     for cookie in nbesf_data.cookies:
         try:
-            new_cookie = DBCookie(
-                site_name=cookie.site_name,
-                content=cookie.content,
-                cookie_name=cookie.cookie_name,
-                cd_milliseconds=cookie.cd_milliseconds,
-                is_universal=cookie.is_universal,
-                tags=cookie.tags,
-            )
+            new_cookie = DBCookie(**model_dump(cookie, exclude={"targets"}))
             cookie_id = await config.add_cookie(new_cookie)
             for target in cookie.targets:
                 await config.add_cookie_target(target.target, target.platform_name, cookie_id)
diff --git a/nonebot_bison/config/subs_io/subs_io.py b/nonebot_bison/config/subs_io/subs_io.py
index 779b6c9..1d8a204 100644
--- a/nonebot_bison/config/subs_io/subs_io.py
+++ b/nonebot_bison/config/subs_io/subs_io.py
@@ -10,7 +10,7 @@ from nonebot.compat import type_validate_python
 from nonebot_plugin_datastore.db import create_session
 from sqlalchemy.orm.strategy_options import selectinload
 
-from .utils import NBESFVerMatchErr
+from .utils import NBESFVerMatchErr, row2dict
 from .nbesf_model import NBESFBase, v1, v2, v3
 from ..db_model import User, Cookie, Target, Subscribe, CookieTarget
 
@@ -59,7 +59,7 @@ async def subscribes_export(selector: Callable[[Select], Select]) -> v3.SubGroup
         )
         cookie_target_data = await sess.scalars(cookie_target_stmt)
 
-    cookie_target_dict = defaultdict(list)
+    cookie_target_dict: dict[Cookie, list[v3.Target]] = defaultdict(list)
     for cookie_target in cookie_target_data:
         target_payload = type_validate_python(v3.Target, cookie_target.target)
         cookie_target_dict[cookie_target.cookie].append(target_payload)
@@ -67,17 +67,10 @@ async def subscribes_export(selector: Callable[[Select], Select]) -> v3.SubGroup
     cookies: list[v3.Cookie] = []
     for cookie, targets in cookie_target_dict.items():
         assert isinstance(cookie, Cookie)
-        cookies.append(
-            v3.Cookie(
-                site_name=cookie.site_name,
-                content=cookie.content,
-                cookie_name=cookie.cookie_name,
-                cd_milliseconds=cookie.cd_milliseconds,
-                is_universal=cookie.is_universal,
-                tags=cookie.tags,
-                targets=targets,
-            )
-        )
+        cookie_dict = row2dict(cookie)
+        cookie_dict["tags"] = cookie.tags
+        cookie_dict["targets"] = targets
+        cookies.append(v3.Cookie(**cookie_dict))
 
     sub_group = v3.SubGroup(groups=groups, cookies=cookies)
 
diff --git a/nonebot_bison/config/subs_io/utils.py b/nonebot_bison/config/subs_io/utils.py
index 181769a..c862b78 100644
--- a/nonebot_bison/config/subs_io/utils.py
+++ b/nonebot_bison/config/subs_io/utils.py
@@ -2,3 +2,11 @@ class NBESFVerMatchErr(Exception): ...
 
 
 class NBESFParseErr(Exception): ...
+
+
+def row2dict(row):
+    d = {}
+    for column in row.__table__.columns:
+        d[column.name] = str(getattr(row, column.name))
+
+    return d