From ff770940ca50517501f8f0036ea2ddcd6773bcaa Mon Sep 17 00:00:00 2001 From: Tilian B <16045812+tilianb@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:38:53 +1000 Subject: [PATCH] Add torrent post-import category action (#1148) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - add a **Change Category** torrent completion action and conditionally show its post-import category/label setting - update kept torrents only after a successful library import - support qBittorrent categories, Transmission labels, Deluge's Label plugin, and rTorrent's `custom1` label - preserve existing Keep/Remove behavior and document the new environment setting ## Behavior Category changes happen from `post_process_cleanup`, after output transfer and post-processing complete. This keeps the existing hardlink flow unchanged. An empty post-import category is a no-op, and client API failures are logged without turning a successful library import into a failure. ## Validation - `pytest -n 0 tests/prowlarr/test_qbittorrent_client.py tests/prowlarr/test_transmission_client.py tests/prowlarr/test_deluge_client.py tests/prowlarr/test_rtorrent_client.py tests/prowlarr/test_handler.py tests/config/test_generate_env_docs.py` — 161 passed - `ruff check` on all changed Python files — passed - `basedpyright` on the changed client implementation files — passed - full `basedpyright shelfmark/download/clients` currently reports two pre-existing errors in the new debrid connection-test code at `settings.py:546` and `settings.py:563`, outside this PR's diff - multi-architecture Docker images built successfully for `linux/amd64` and `linux/arm64` --- docs/environment-variables.md | 16 +++++-- shelfmark/download/clients/__init__.py | 13 +++++ shelfmark/download/clients/base_handler.py | 34 +++++++++++-- shelfmark/download/clients/deluge.py | 16 ++++++- shelfmark/download/clients/qbittorrent.py | 20 ++++++++ shelfmark/download/clients/rtorrent.py | 18 +++++++ shelfmark/download/clients/settings.py | 11 ++++- shelfmark/download/clients/transmission.py | 11 +++++ tests/prowlarr/test_deluge_client.py | 11 +++++ tests/prowlarr/test_handler.py | 55 ++++++++++++++++++++++ tests/prowlarr/test_qbittorrent_client.py | 14 ++++++ tests/prowlarr/test_rtorrent_client.py | 10 ++++ tests/prowlarr/test_transmission_client.py | 13 +++++ 13 files changed, 233 insertions(+), 9 deletions(-) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 828308a..cb01ec9 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -1534,7 +1534,8 @@ How long to keep cached search results before they expire. | `RTORRENT_LABEL` | Label to assign to ebook downloads in rTorrent | string | `cwabd` | | `RTORRENT_AUDIOBOOK_LABEL` | Label to assign to audiobook downloads in rTorrent (falls back to Book Label if not set) | string | _none_ | | `RTORRENT_DOWNLOAD_DIR` | Server-side directory where torrents are downloaded (optional, uses rTorrent default if not specified) | string | _none_ | -| `PROWLARR_TORRENT_ACTION` | Remove deletes the torrent from your client immediately after import (stops seeding, files are kept); Keep leaves it in the client to continue seeding | string (choice) | `keep` | +| `PROWLARR_TORRENT_ACTION` | Choose whether to keep, remove, or move the torrent to another category or label after import | string (choice) | `keep` | +| `PROWLARR_TORRENT_POST_IMPORT_CATEGORY` | Category or label to assign after a successful import | string | _empty string_ | | `PROWLARR_USENET_CLIENT` | Choose which usenet client to use | string (choice) | _empty string_ | | `NZBGET_URL` | URL of your NZBGet instance | string | _none_ | | `NZBGET_USERNAME` | NZBGet control username | string | `nzbget` | @@ -1798,11 +1799,20 @@ Server-side directory where torrents are downloaded (optional, uses rTorrent def **Torrent Completion Action** -Remove deletes the torrent from your client immediately after import (stops seeding, files are kept); Keep leaves it in the client to continue seeding +Choose whether to keep, remove, or move the torrent to another category or label after import - **Type:** string (choice) - **Default:** `keep` -- **Options:** `keep` (Keep), `remove` (Remove) +- **Options:** `keep` (Keep), `remove` (Remove), `change_category` (Change Category) + +#### `PROWLARR_TORRENT_POST_IMPORT_CATEGORY` + +**Post-Import Category** + +Category or label to assign after a successful import + +- **Type:** string +- **Default:** _empty string_ #### `PROWLARR_USENET_CLIENT` diff --git a/shelfmark/download/clients/__init__.py b/shelfmark/download/clients/__init__.py index 57d84d6..c0b7f25 100644 --- a/shelfmark/download/clients/__init__.py +++ b/shelfmark/download/clients/__init__.py @@ -325,6 +325,19 @@ class DownloadClient(ABC): """ + def set_category(self, download_id: str, category: str) -> bool: + """Update a download's category or label when supported by the client. + + Args: + download_id: The client-specific download ID. + category: Category or label to assign. + + Returns: + True if the category was updated, otherwise False. + + """ + return False + @abstractmethod def get_download_path(self, download_id: str) -> str | None: """Get the path where files were downloaded. diff --git a/shelfmark/download/clients/base_handler.py b/shelfmark/download/clients/base_handler.py index 4fb5893..4f76d3b 100644 --- a/shelfmark/download/clients/base_handler.py +++ b/shelfmark/download/clients/base_handler.py @@ -284,17 +284,45 @@ class ExternalClientHandler(DownloadHandler, ABC): ) elif protocol == "torrent": - if config.get("PROWLARR_TORRENT_ACTION", "keep") != "remove": + torrent_action = config.get("PROWLARR_TORRENT_ACTION", "keep") + if torrent_action == "remove": + try: + client.remove(download_id, delete_files=False) + except _CLIENT_CLEANUP_ERRORS as e: + logger.warning( + "Failed to remove torrent %s from %s: %s", + download_id, + getattr(client, "name", "client"), + e, + ) return + + if torrent_action != "change_category": + return + + post_import_category = normalize_optional_text( + config.get("PROWLARR_TORRENT_POST_IMPORT_CATEGORY", "") + ) + if post_import_category is None: + return + try: - client.remove(download_id, delete_files=False) + category_updated = client.set_category(download_id, post_import_category) except _CLIENT_CLEANUP_ERRORS as e: logger.warning( - "Failed to remove torrent %s from %s: %s", + "Failed to set post-import category for torrent %s in %s: %s", download_id, getattr(client, "name", "client"), e, ) + return + + if not category_updated: + logger.warning( + "Failed to set post-import category for torrent %s in %s", + download_id, + getattr(client, "name", "client"), + ) def _remove_usenet_download( self, diff --git a/shelfmark/download/clients/deluge.py b/shelfmark/download/clients/deluge.py index 6b0620d..e55104a 100644 --- a/shelfmark/download/clients/deluge.py +++ b/shelfmark/download/clients/deluge.py @@ -227,10 +227,10 @@ class DelugeClient(DownloadClient): return self._rpc_call("daemon.info") - def _try_set_label(self, torrent_id: str, label: str) -> None: + def _try_set_label(self, torrent_id: str, label: str) -> bool: """Best-effort label assignment (requires Deluge Label plugin).""" if not label: - return + return False try: # label.add will error if the plugin is unavailable or the label exists. @@ -240,6 +240,9 @@ class DelugeClient(DownloadClient): self._rpc_call("label.set_torrent", torrent_id, label) except _DELUGE_CLIENT_ERRORS as e: logger.debug("Could not set Deluge label '%s' for %s: %s", label, torrent_id, e) + return False + else: + return True @staticmethod def is_configured() -> bool: @@ -422,6 +425,15 @@ class DelugeClient(DownloadClient): else: return False + def set_category(self, download_id: str, category: str) -> bool: + """Assign a label to a torrent using Deluge's Label plugin.""" + try: + self._ensure_connected() + return self._try_set_label(download_id, category) + except _DELUGE_CLIENT_ERRORS as e: + self._log_error("set_category", e) + return False + def get_download_path(self, download_id: str) -> str | None: """Return the resolved download path for a Deluge torrent.""" try: diff --git a/shelfmark/download/clients/qbittorrent.py b/shelfmark/download/clients/qbittorrent.py index 431b755..d150a71 100644 --- a/shelfmark/download/clients/qbittorrent.py +++ b/shelfmark/download/clients/qbittorrent.py @@ -673,6 +673,26 @@ class QBittorrentClient(DownloadClient): else: return True + def set_category(self, download_id: str, category: str) -> bool: + """Assign a category to a torrent in qBittorrent.""" + try: + try: + self._client.torrents_create_category(name=category) + except _QBITTORRENT_CLIENT_ERRORS as e: + if "Conflict" not in type(e).__name__ and "409" not in str(e): + logger.debug("Could not create category '%s': %s", category, e) + + self._client.torrents_set_category( + torrent_hashes=download_id, + category=category, + ) + logger.info("Set qBittorrent category for %s to '%s'", download_id, category) + except _QBITTORRENT_CLIENT_ERRORS as e: + self._log_error("set_category", e) + return False + else: + return True + def get_download_path(self, download_id: str) -> str | None: """Get the path where torrent files are located. diff --git a/shelfmark/download/clients/rtorrent.py b/shelfmark/download/clients/rtorrent.py index 457376a..7bf3d55 100644 --- a/shelfmark/download/clients/rtorrent.py +++ b/shelfmark/download/clients/rtorrent.py @@ -47,7 +47,13 @@ class _RTorrentLoadProtocol(Protocol): def start(self, target: str, url: str, commands: str) -> object: ... +class _RTorrentCustom1Protocol(Protocol): + def set(self, download_id: str, value: str) -> object: ... + + class _RTorrentDownloadProtocol(Protocol): + custom1: _RTorrentCustom1Protocol + def multicall2(self, *args: object) -> list[list[Any]]: ... def delete_tied(self, download_id: str) -> object: ... @@ -351,6 +357,18 @@ class RTorrentClient(DownloadClient): else: return True + def set_category(self, download_id: str, category: str) -> bool: + """Assign a label to a torrent using rTorrent's custom1 field.""" + try: + self._rpc.d.custom1.set(download_id, category) + logger.info("Set rTorrent label for %s to '%s'", download_id, category) + except _RTORRENT_CLIENT_ERRORS as e: + error_type = type(e).__name__ + logger.exception("rTorrent set_category failed (%s)", error_type) + return False + else: + return True + def get_download_path(self, download_id: str) -> str | None: """Get the path where torrent files are located. diff --git a/shelfmark/download/clients/settings.py b/shelfmark/download/clients/settings.py index 197f05a..6e2a04d 100644 --- a/shelfmark/download/clients/settings.py +++ b/shelfmark/download/clients/settings.py @@ -854,14 +854,23 @@ def prowlarr_clients_settings() -> list[SettingsField]: SelectField( key="PROWLARR_TORRENT_ACTION", label="Torrent Completion Action", - description="Remove deletes the torrent from your client immediately after import (stops seeding, files are kept); Keep leaves it in the client to continue seeding", + description="Choose whether to keep, remove, or move the torrent to another category or label after import", options=[ {"value": "keep", "label": "Keep"}, {"value": "remove", "label": "Remove"}, + {"value": "change_category", "label": "Change Category"}, ], default="keep", show_when={"field": "PROWLARR_TORRENT_CLIENT", "notEmpty": True}, ), + TextField( + key="PROWLARR_TORRENT_POST_IMPORT_CATEGORY", + label="Post-Import Category", + description="Category or label to assign after a successful import", + placeholder="imported", + default="", + show_when={"field": "PROWLARR_TORRENT_ACTION", "value": "change_category"}, + ), # --- Usenet Client Selection --- HeadingField( key="usenet_heading", diff --git a/shelfmark/download/clients/transmission.py b/shelfmark/download/clients/transmission.py index 9bd6a5d..42d5fd1 100644 --- a/shelfmark/download/clients/transmission.py +++ b/shelfmark/download/clients/transmission.py @@ -390,6 +390,17 @@ class TransmissionClient(DownloadClient): else: return True + def set_category(self, download_id: str, category: str) -> bool: + """Replace a torrent's labels with the post-import label.""" + try: + self._client.change_torrent(ids=download_id, labels=[category]) + logger.info("Set Transmission label for %s to '%s'", download_id, category) + except _TRANSMISSION_CLIENT_ERRORS as e: + self._log_error("set_category", e) + return False + else: + return True + def get_download_path(self, download_id: str) -> str | None: """Get the path where torrent files are located. diff --git a/tests/prowlarr/test_deluge_client.py b/tests/prowlarr/test_deluge_client.py index 563ad74..4002e9f 100644 --- a/tests/prowlarr/test_deluge_client.py +++ b/tests/prowlarr/test_deluge_client.py @@ -15,6 +15,17 @@ def make_config_getter(values): return getter +def test_set_category_uses_label_plugin(monkeypatch): + from shelfmark.download.clients.deluge import DelugeClient + + client = DelugeClient.__new__(DelugeClient) + monkeypatch.setattr(client, "_ensure_connected", MagicMock()) + monkeypatch.setattr(client, "_try_set_label", MagicMock(return_value=True)) + + assert client.set_category("abc123", "imported") is True + client._try_set_label.assert_called_once_with("abc123", "imported") + + class TestDelugeClientAddDownload: """Tests for DelugeClient.add_download().""" diff --git a/tests/prowlarr/test_handler.py b/tests/prowlarr/test_handler.py index a3201b2..62466b0 100644 --- a/tests/prowlarr/test_handler.py +++ b/tests/prowlarr/test_handler.py @@ -1456,6 +1456,61 @@ class TestProwlarrHandlerFileStaging: class TestProwlarrHandlerPostProcessCleanup: + def test_torrent_change_category_sets_post_import_category(self): + handler = ProwlarrHandler() + task = DownloadTask(task_id="torrent-category", source="prowlarr", title="Test") + + mock_client = MagicMock() + mock_client.name = "qbittorrent" + mock_client.set_category.return_value = True + handler._cleanup_refs[task.task_id] = (mock_client, "abc123", "torrent") + + config_values = { + "PROWLARR_TORRENT_ACTION": "change_category", + "PROWLARR_TORRENT_POST_IMPORT_CATEGORY": "imported", + } + with patch( + "shelfmark.download.clients.base_handler.config.get", + side_effect=lambda key, default="": config_values.get(key, default), + ): + handler.post_process_cleanup(task, success=True) + + mock_client.set_category.assert_called_once_with("abc123", "imported") + mock_client.remove.assert_not_called() + + def test_torrent_keep_does_not_change_category(self): + handler = ProwlarrHandler() + task = DownloadTask(task_id="torrent-no-category", source="prowlarr", title="Test") + + mock_client = MagicMock() + handler._cleanup_refs[task.task_id] = (mock_client, "abc123", "torrent") + + with patch("shelfmark.download.clients.base_handler.config.get", return_value="keep"): + handler.post_process_cleanup(task, success=True) + + mock_client.set_category.assert_not_called() + mock_client.remove.assert_not_called() + + def test_torrent_change_category_ignores_empty_category(self): + handler = ProwlarrHandler() + task = DownloadTask(task_id="torrent-empty-category", source="prowlarr", title="Test") + + mock_client = MagicMock() + handler._cleanup_refs[task.task_id] = (mock_client, "abc123", "torrent") + + config_values = { + "PROWLARR_TORRENT_ACTION": "change_category", + "PROWLARR_TORRENT_POST_IMPORT_CATEGORY": "", + } + with patch( + "shelfmark.download.clients.base_handler.config.get", + side_effect=lambda key, default="": config_values.get(key, default), + ): + handler.post_process_cleanup(task, success=True) + + mock_client.set_category.assert_not_called() + mock_client.remove.assert_not_called() + def test_usenet_move_triggers_client_cleanup(self): handler = ProwlarrHandler() task = DownloadTask(task_id="cleanup-test", source="prowlarr", title="Test") diff --git a/tests/prowlarr/test_qbittorrent_client.py b/tests/prowlarr/test_qbittorrent_client.py index 540fb52..f0465b8 100644 --- a/tests/prowlarr/test_qbittorrent_client.py +++ b/tests/prowlarr/test_qbittorrent_client.py @@ -59,6 +59,20 @@ def create_mock_session_response(torrents, status_code=200): return mock_response +def test_set_category_creates_and_assigns_category(): + from shelfmark.download.clients.qbittorrent import QBittorrentClient + + client = QBittorrentClient.__new__(QBittorrentClient) + client._client = MagicMock() + + assert client.set_category("abc123", "imported") is True + client._client.torrents_create_category.assert_called_once_with(name="imported") + client._client.torrents_set_category.assert_called_once_with( + torrent_hashes="abc123", + category="imported", + ) + + class TestQBittorrentClientIsConfigured: """Tests for QBittorrentClient.is_configured().""" diff --git a/tests/prowlarr/test_rtorrent_client.py b/tests/prowlarr/test_rtorrent_client.py index 52ec424..a681491 100644 --- a/tests/prowlarr/test_rtorrent_client.py +++ b/tests/prowlarr/test_rtorrent_client.py @@ -27,6 +27,16 @@ def create_mock_xmlrpc_module(): return mock_module +def test_set_category_updates_custom1_label(): + from shelfmark.download.clients.rtorrent import RTorrentClient + + client = RTorrentClient.__new__(RTorrentClient) + client._rpc = MagicMock() + + assert client.set_category("abc123", "imported") is True + client._rpc.d.custom1.set.assert_called_once_with("abc123", "imported") + + class TestRTorrentClientIsConfigured: """Tests for RTorrentClient.is_configured().""" diff --git a/tests/prowlarr/test_transmission_client.py b/tests/prowlarr/test_transmission_client.py index 3a80fda..7372c42 100644 --- a/tests/prowlarr/test_transmission_client.py +++ b/tests/prowlarr/test_transmission_client.py @@ -68,6 +68,19 @@ def create_mock_transmission_rpc_module(): return mock_module +def test_set_category_replaces_labels(): + from shelfmark.download.clients.transmission import TransmissionClient + + client = TransmissionClient.__new__(TransmissionClient) + client._client = MagicMock() + + assert client.set_category("abc123", "imported") is True + client._client.change_torrent.assert_called_once_with( + ids="abc123", + labels=["imported"], + ) + + class TestTransmissionClientIsConfigured: """Tests for TransmissionClient.is_configured()."""