From 1425b153336d318296bfcc5bf04e3f70e0037cd1 Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Mon, 27 May 2024 11:16:41 +0800
Subject: [PATCH] fix #2192

---
 app/chain/__init__.py                    |  2 +-
 app/modules/qbittorrent/__init__.py      |  2 +-
 app/modules/transmission/__init__.py     | 11 ++++++--
 app/modules/transmission/transmission.py | 33 ++++++++++++++++++------
 4 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/app/chain/__init__.py b/app/chain/__init__.py
index 6b07e91c..873030cd 100644
--- a/app/chain/__init__.py
+++ b/app/chain/__init__.py
@@ -375,7 +375,7 @@ class ChainBase(metaclass=ABCMeta):
                                transfer_type=transfer_type, target=target, episodes_info=episodes_info,
                                scrape=scrape)
 
-    def transfer_completed(self, hashs: Union[str, list], path: Path = None,
+    def transfer_completed(self, hashs: str, path: Path = None,
                            downloader: str = settings.DEFAULT_DOWNLOADER) -> None:
         """
         转移完成后的处理
diff --git a/app/modules/qbittorrent/__init__.py b/app/modules/qbittorrent/__init__.py
index 28203da8..ac7cdae3 100644
--- a/app/modules/qbittorrent/__init__.py
+++ b/app/modules/qbittorrent/__init__.py
@@ -242,7 +242,7 @@ class QbittorrentModule(_ModuleBase):
             return None
         return ret_torrents
 
-    def transfer_completed(self, hashs: Union[str, list], path: Path = None,
+    def transfer_completed(self, hashs: str, path: Path = None,
                            downloader: str = settings.DEFAULT_DOWNLOADER) -> None:
         """
         转移完成后的处理
diff --git a/app/modules/transmission/__init__.py b/app/modules/transmission/__init__.py
index 91bb3ef7..bb8f79fc 100644
--- a/app/modules/transmission/__init__.py
+++ b/app/modules/transmission/__init__.py
@@ -230,7 +230,7 @@ class TransmissionModule(_ModuleBase):
             return None
         return ret_torrents
 
-    def transfer_completed(self, hashs: Union[str, list], path: Path = None,
+    def transfer_completed(self, hashs: str, path: Path = None,
                            downloader: str = settings.DEFAULT_DOWNLOADER) -> None:
         """
         转移完成后的处理
@@ -241,7 +241,14 @@ class TransmissionModule(_ModuleBase):
         """
         if downloader != "transmission":
             return None
-        self.transmission.set_torrent_tag(ids=hashs, tags=['已整理'])
+        # 获取原标签
+        org_tags = self.transmission.get_torrent_tags(ids=hashs)
+        # 种子打上已整理标签
+        if org_tags:
+            tags = org_tags + ['已整理']
+        else:
+            tags = ['已整理']
+        self.transmission.set_torrent_tag(ids=hashs, tags=tags)
         # 移动模式删除种子
         if settings.TRANSFER_TYPE == "move":
             if self.remove_torrents(hashs):
diff --git a/app/modules/transmission/transmission.py b/app/modules/transmission/transmission.py
index cd40fc24..87447a3e 100644
--- a/app/modules/transmission/transmission.py
+++ b/app/modules/transmission/transmission.py
@@ -2,7 +2,7 @@ from typing import Optional, Union, Tuple, List, Dict
 
 import transmission_rpc
 from transmission_rpc import Client, Torrent, File
-from transmission_rpc.session import SessionStats
+from transmission_rpc.session import SessionStats, Session
 
 from app.core.config import settings
 from app.log import logger
@@ -130,21 +130,38 @@ class Transmission:
             logger.error(f"获取正在下载的种子列表出错：{str(err)}")
             return None
 
-    def set_torrent_tag(self, ids: str, tags: list) -> bool:
+    def set_torrent_tag(self, ids: str, tags: list, org_tags: list = None) -> bool:
         """
-        设置种子标签
+        设置种子标签，注意TR默认会覆盖原有标签，如需追加需传入原有标签
         """
         if not self.trc:
             return False
         if not ids or not tags:
             return False
         try:
-            self.trc.change_torrent(labels=tags, ids=ids)
+            self.trc.change_torrent(labels=list(set((org_tags or []) + tags)), ids=ids)
             return True
         except Exception as err:
             logger.error(f"设置种子标签出错：{str(err)}")
             return False
 
+    def get_torrent_tags(self, ids: str) -> List[str]:
+        """
+        获取所有种子标签
+        """
+        if not self.trc:
+            return []
+        try:
+            torrent = self.trc.get_torrents(ids=ids, arguments=self._trarg)
+            if torrent:
+                labels = [str(tag).strip()
+                          for tag in torrent.labels] if hasattr(torrent, "labels") else []
+                return labels
+        except Exception as err:
+            logger.error(f"获取种子标签出错：{str(err)}")
+            return []
+        return []
+
     def add_torrent(self, content: Union[str, bytes],
                     is_paused: bool = False,
                     download_dir: str = None,
@@ -397,15 +414,15 @@ class Transmission:
             logger.error(f"修改tracker出错：{str(err)}")
             return False
 
-    def get_session(self) -> Dict[str, Union[int, bool, str]]:
+    def get_session(self) -> Optional[Session]:
         """
         获取Transmission当前的会话信息和配置设置
-        :return dict or False
+        :return dict
         """
         if not self.trc:
-            return False
+            return None
         try:
             return self.trc.get_session()
         except Exception as err:
             logger.error(f"获取session出错：{str(err)}")
-            return False
+            return None