diff --git a/pywxdump/db/dbFavorite.py b/pywxdump/db/dbFavorite.py index 568c128..677c64f 100644 --- a/pywxdump/db/dbFavorite.py +++ b/pywxdump/db/dbFavorite.py @@ -5,6 +5,8 @@ # Author: xaoyaoo # Date: 2024/05/18 # ------------------------------------------------------------------------------- +from collections import defaultdict + from .dbbase import DatabaseBase from .utils import timestamp2str, xml2dict @@ -133,7 +135,7 @@ class FavoriteHandler(DatabaseBase): pf.columns = FavItemsFields.keys() # set column names pf["UpdateTime"] = pf["UpdateTime"].apply(timestamp2str) # 处理时间 pf["XmlBuf"] = pf["XmlBuf"].apply(xml2dict) # 处理xml - pf["TypeName"] = pf["Type"].apply(FavoriteTypeId2Name) # 添加类型名称列 + pf["TypeName"] = pf["Type"].apply(Favorite_type_converter) # 添加类型名称列 pf["FavData"] = pf["FavLocalID"].apply(lambda x: FavDataDict.get(x, [])) # 添加数据列 pf["Tags"] = pf["FavLocalID"].apply(lambda x: FavTagsDict.get(x, [])) # 添加标签列 pf = pf.fillna("") # 去掉Nan @@ -141,8 +143,15 @@ class FavoriteHandler(DatabaseBase): return rdata -def FavoriteTypeId2Name(Type): - TypeNameDict = { +def Favorite_type_converter(type_id_or_name: [str, int]): + """ + 收藏类型ID与名称转换 + 名称(str)=>ID(int) + ID(int)=>名称(str) + :param type_id_or_name: 消息类型ID或名称 + :return: 消息类型名称或ID + """ + type_name_dict = defaultdict(lambda: "未知", { 1: "文本", # 文本 已测试 2: "图片", # 图片 已测试 3: "语音", # 语音 @@ -154,5 +163,11 @@ def FavoriteTypeId2Name(Type): 14: "聊天记录", # 聊天记录 已测试 16: "群聊视频", # 群聊中的视频 可能 18: "笔记" # 笔记 已测试 - } - return TypeNameDict.get(Type, "未知") + }) + + if isinstance(type_id_or_name, int): + return type_name_dict[type_id_or_name] + elif isinstance(type_id_or_name, str): + return next((k for k, v in type_name_dict.items() if v == type_id_or_name), (0, 0)) + else: + raise ValueError("Invalid input type")