diff --git a/stig/client/aiotransmission/api_torrent.py b/stig/client/aiotransmission/api_torrent.py index 097ecc52..ba9f5834 100644 --- a/stig/client/aiotransmission/api_torrent.py +++ b/stig/client/aiotransmission/api_torrent.py @@ -14,6 +14,7 @@ import time from collections import abc from string import hexdigits as HEXDIGITS +from enum import Enum from natsort import humansorted @@ -132,13 +133,14 @@ async def _abs_download_path(self, path): return Response(success=True, path=SmartCmpPath(abs_path)) return Response(success=True, path=SmartCmpPath(path)) - async def add(self, torrent, stopped=False, path=None): + async def add(self, torrent, stopped=False, path=None, labels=[]): """ Add torrent from file, URL or hash torrent: Path to local file, web/magnet link or hash stopped: False to start downloading immediately, True otherwise path: Download directory or `None` for default directory + labels: List of labels Return Response with the following properties: torrent: Torrent object with the keys 'id' and 'name' if the @@ -150,7 +152,7 @@ async def add(self, torrent, stopped=False, path=None): errors: List of error messages """ torrent_str = torrent - args = {'paused': bool(stopped)} + args = {'paused': bool(stopped), 'labels': labels} if path is not None: response = await self._abs_download_path(path) @@ -209,8 +211,17 @@ async def add(self, torrent, stopped=False, path=None): success = False elif 'torrent-added' in result: info = result['torrent-added'] - msgs = ('Added %s' % info['name'],) + msgs = ['Added %s' % info['name']] success = True + # Before rpc version 17 torrent-add did not take the labels + # argument, so we have to send a follow-up request + if labels and self.rpc.rpcversion < 17: + response = await self.labels_add((info['id'],), labels) + success = response.success + if response.success: + msgs.append('Labeled %s with %s' % (info['name'], ', '.join(labels))) + else: + errors = ('Could not label added torrents: %s' % response.errors,) else: raise RuntimeError('Malformed response: %r' % (result,)) torrent = Torrent({'id': info['id'], 'name': info['name']}) @@ -1099,3 +1110,135 @@ def check(t): return await self._torrent_action(self.rpc.torrent_reannounce, torrents, check=check, check_keys=('status', 'trackers', 'time-manual-announce-allowed')) + + label_manage_mode = Enum('label_manage_mode', 'ADD SET REMOVE') + + async def labels_add(self, torrents, labels): + """ + Add labels(s) to torrents + + See '_labels_manage' method + """ + return await self._labels_manage(torrents, labels, self.label_manage_mode.ADD) + + async def labels_remove(self, torrents, labels): + """ + Add labels(s) to torrents + + See '_labels_manage' method + """ + return await self._labels_manage(torrents, labels, + self.label_manage_mode.REMOVE) + + async def labels_set(self, torrents, labels): + """ + Add labels(s) to torrents + + See '_labels_manage' method + """ + return await self._labels_manage(torrents, labels, self.label_manage_mode.SET) + + async def labels_clear(self, torrents): + """ + Clear labels(s) from torrents + + torrents: See `torrents` method + + Return Response with the following properties: + torrents: Tuple of Torrents with the keys 'id' and 'name' + success: True if all RPC requests returned successfully + msgs: List of info messages + errors: List of error messages + """ + + msgs = [] + args = {'labels': []} + response = await self._torrent_action(self.rpc.torrent_set, torrents, method_args=args) + if not response.success: + return Response(success=False, torrents=(), msgs=[], errors=response.errors) + else: + for t in response.torrents: + msgs.append('%s: Cleared labels' % t['name']) + return Response(success=True, torrents=response.torrents, + msgs=msgs, errors=[]) + + async def _labels_manage(self, torrents, labels, mode): + """ + Set/add/remove torrent labels(s) + + torrents: See `torrents` method + labels: Iterable of labels + mode: set/add/remove + + Return Response with the following properties: + torrents: Tuple of Torrents with the keys 'id', 'name' and 'labels' + success: True if any labels were added/removed/set, False if no labels were + added or if new label lists could not be retrieved + msgs: List of info messages + errors: List of error messages + """ + if not labels: + return Response(success=False, torrents=(), errors=('No labels given',)) + labels = set(labels) + + # Transmission only allows *setting* the label list, so first we get + # any existing labels + response = await self.torrents(torrents, keys=('id', 'name', 'labels',)) + if not response.success: + return Response(success=False, torrents=(), errors=response.errors) + + # Map torrent IDs to that torrent's current labels + tor_dict = {t['id']:t for t in response.torrents} + set_funcs = { + self.label_manage_mode.ADD: lambda x, y: x.union(y), + self.label_manage_mode.REMOVE: lambda x, y: x.difference(y), + self.label_manage_mode.SET: lambda x, y: y, + } + label_dict = { + t['id']: frozenset(set_funcs[mode](t['labels'], labels)) + for t in response.torrents + } + # Collating by label set reduces the number of requests we need to send + # to one per label set + inv_label_dict = {} + for tid, ls in label_dict.items(): + inv_label_dict.setdefault(ls, []).append(tid) + + # Add/remove/set labels + modded_any = False + msgs = [] + errors = [] + verbs = { + self.label_manage_mode.ADD: 'Adding', + self.label_manage_mode.REMOVE: 'Removing', + self.label_manage_mode.SET: 'Setting', + } + for ls, tids in inv_label_dict.items(): + args = {'labels': list(ls)} + response = await self._torrent_action(self.rpc.torrent_set, tids, method_args=args) + if not response.success: + errors.extend(response.errors) + continue + for t in tids: + if mode == self.label_manage_mode.ADD: + mod_labels = labels.difference(tor_dict[t]['labels']) + elif mode == self.label_manage_mode.REMOVE: + mod_labels = tor_dict[t]['labels'].intersection(labels) + elif mode == self.label_manage_mode.SET: + mod_labels = tor_dict[t]['labels'].symmetric_difference(labels) + if mod_labels: + if mode == self.label_manage_mode.SET: + mod_labels = labels + msgs.append('%s: %s labels: %s' % ( + tor_dict[t]['name'], verbs[mode], ', '.join(mod_labels)) + ) + modded_any = True + + response = await self.torrents(torrents, keys=('id', 'name', 'labels',)) + if not modded_any: + errors.append('No labels were changed') + if not response.success: + errors.append('Failed to read new labels from transmission') + + return Response(success=response.success and modded_any, + torrents=response.torrents, msgs=msgs, errors=errors) diff --git a/stig/client/aiotransmission/torrent.py b/stig/client/aiotransmission/torrent.py index 9bce9313..9a12b888 100644 --- a/stig/client/aiotransmission/torrent.py +++ b/stig/client/aiotransmission/torrent.py @@ -429,6 +429,8 @@ def __new__(cls, raw_torrent): 'trackers' : ('trackerStats', 'name', 'id'), 'peers' : ('peers', 'totalSize', 'name'), 'files' : ('files', 'fileStats', 'downloadDir'), + + 'labels' : ('labels',), } @@ -577,7 +579,7 @@ class TorrentFields(tuple): 'downloadedEver', 'downloadLimit', 'downloadLimited', 'downloadLimitMode', 'error', 'errorString', 'eta', 'etaIdle', 'hashString', 'haveUnchecked', 'haveValid', 'honorsSessionLimits', - 'id', 'isFinished', 'isPrivate', 'isStalled', 'lastAnnounceTime', + 'id', 'isFinished', 'isPrivate', 'isStalled', 'labels', 'lastAnnounceTime', 'lastScrapeTime', 'leftUntilDone', 'magnetLink', 'manualAnnounceTime', 'maxConnectedPeers', 'metadataPercentComplete', 'name', 'nextAnnounceTime', diff --git a/stig/client/base.py b/stig/client/base.py index edfa1328..5c81de9b 100644 --- a/stig/client/base.py +++ b/stig/client/base.py @@ -82,6 +82,8 @@ class TorrentBase(abc.Mapping): 'trackers' : tuple, 'peers' : tuple, 'files' : None, + + 'labels' : set, } def update(self, raw_torrent): diff --git a/stig/client/filters/torrent.py b/stig/client/filters/torrent.py index 9f47a1cd..b4f5d5b4 100644 --- a/stig/client/filters/torrent.py +++ b/stig/client/filters/torrent.py @@ -202,6 +202,13 @@ class _SingleFilter(Filter): needed_keys=('trackers',), aliases=('trk',), description=_desc('... domain of the announce URL of trackers')), + 'label' : CmpFilterSpec(value_getter=lambda t: t['labels'], + value_matcher=lambda t, op, v: + any(op(lbl, v) for lbl in t['labels']), + value_type=str, + needed_keys=('labels',), + aliases=('lbl',), + description=_desc('... labels')), 'eta' : CmpFilterSpec(value_getter=lambda t: t['timespan-eta'], value_matcher=lambda t, op, v: cmp_timestamp_or_timdelta(t['timespan-eta'], op, v), diff --git a/stig/commands/base/torrent.py b/stig/commands/base/torrent.py index c57289f4..3c2d7ea1 100644 --- a/stig/commands/base/torrent.py +++ b/stig/commands/base/torrent.py @@ -32,7 +32,8 @@ class AddTorrentsCmdbase(metaclass=CommandMeta): description = 'Download torrents' usage = ('add [] ...',) examples = ('add 72d7a3179da3de7a76b98f3782c31843e3f818ee', - 'add --stopped http://example.org/something.torrent') + 'add --stopped http://example.org/something.torrent', + 'add --labels linux,iso https://archlinux.org/releng/releases/2022.04.05/torrent/') argspecs = ( {'names': ('TORRENT',), 'nargs': '+', 'description': 'Link or path to torrent file, magnet link or info hash'}, @@ -43,16 +44,22 @@ class AddTorrentsCmdbase(metaclass=CommandMeta): {'names': ('--path','-p'), 'description': ('Custom download directory for added torrent(s) ' 'relative to "srv.path.complete" setting')}, + + {'names': ('--labels','-l'), + 'description': 'Comma-separated list of labels'}, ) - async def run(self, TORRENT, stopped, path): + async def run(self, TORRENT, stopped, path, labels): success = True force_torrentlist_update = False + if labels: + labels = labels.split(',') for source in TORRENT: source_abs_path = self.make_path_absolute(source) response = await self.make_request(objects.srvapi.torrent.add(source_abs_path, stopped=stopped, - path=path)) + path=path, + labels=labels)) success = success and response.success force_torrentlist_update = force_torrentlist_update or success @@ -78,6 +85,9 @@ def completion_candidates_params(cls, option, args): return candidates.fs_path(args.curarg.before_cursor, base=objects.cfg['srv.path.complete'], directories_only=True) + if option == '--labels': + curlbl = args.curarg.split(',')[-1] + return candidates.labels(curlbl) class TorrentDetailsCmdbase(mixin.get_single_torrent, metaclass=CommandMeta): @@ -551,3 +561,64 @@ async def run(self, TORRENT_FILTER): def completion_candidates_posargs(cls, args): """Complete positional arguments""" return candidates.torrent_filter(args.curarg) + + +class LabelCmd(metaclass=CommandMeta): + name = 'label' + provides = set() + category = 'torrent' + description = 'Manipulate torrent labels' + usage = ('label [] ... <[LABEL][,LABEL...]>',) + examples = ('label iso,linux id=34', + 'label -r iso,linux id=34', + 'label -c id=34') + argspecs = ( + {'names': ('LABELS',), + 'description': ('Comma-separated list of labels to add/remove'), + 'nargs': '?', 'default': ''}, + make_X_FILTER_spec('TORRENT', or_focused=True, nargs='*'), + {'names': ('--clear','-c'), 'action': 'store_true', + 'description': 'Clear all labels'}, + {'names': ('--remove','-r'), 'action': 'store_true', + 'description': 'Remove labels rather than adding them'}, + {'names': ('--set','-s'), 'action': 'store_true', + 'description': 'Set labels to exactly the LABELS argument', + 'dest': '_set'}, + {'names': ('--quiet','-q'), 'action': 'store_true', + 'description': 'Do not show new label(s)'}, + ) + + async def run(self, LABELS, TORRENT_FILTER, _set, remove, clear, quiet): + if not (_set ^ remove ^ clear) and (_set or remove or clear): + raise CmdError('At most one of --set/s, --remove/r, --clear,-c can be present.') + return + + if clear: + def handler(tfilter, labels): + return objects.srvapi.torrent.labels_clear(tfilter) + elif _set: + handler = objects.srvapi.torrent.labels_set + elif remove: + handler = objects.srvapi.torrent.labels_remove + else: + handler = objects.srvapi.torrent.labels_add + + labels = LABELS.split(',') + + try: + tfilter = self.select_torrents(TORRENT_FILTER, + allow_no_filter=False, + discover_torrent=True) + except ValueError as e: + raise CmdError(e) + + response = await self.make_request(handler(tfilter, labels), + polling_frenzy=True, quiet=quiet) + if not response.success: + raise CmdError() + + @classmethod + def completion_candidates_posargs(cls, args): + """Complete positional arguments""" + curlbl = args.curarg.split(',')[-1] + return candidates.labels(curlbl) diff --git a/stig/commands/cli/torrent.py b/stig/commands/cli/torrent.py index 0c33f8b8..0a00b080 100644 --- a/stig/commands/cli/torrent.py +++ b/stig/commands/cli/torrent.py @@ -167,3 +167,7 @@ class StopTorrentsCmd(base.StopTorrentsCmdbase, class VerifyTorrentsCmd(base.VerifyTorrentsCmdbase, mixin.make_request, mixin.select_torrents): provides = {'cli'} + + +class LabelCmd(base.LabelCmd, mixin.make_request, mixin.select_torrents): + provides = {'cli'} diff --git a/stig/commands/cmdbase.py b/stig/commands/cmdbase.py index 66e71ebb..01958ac4 100644 --- a/stig/commands/cmdbase.py +++ b/stig/commands/cmdbase.py @@ -142,7 +142,7 @@ def __init__(self, args=(), argv=(), info_handler=None, error_handler=None, **kw kwargs = {} for argspec in self.argspecs: # First name is the kwarg for run() - key = argspec['names'][0].lstrip('-').replace('-', '_') + key = argspec.get('dest') or argspec['names'][0].lstrip('-').replace('-', '_') value = getattr(args_parsed, key) kwargs[key.replace(' ', '_')] = value self._args = kwargs diff --git a/stig/commands/tui/torrent.py b/stig/commands/tui/torrent.py index 48efee5b..051a91a0 100644 --- a/stig/commands/tui/torrent.py +++ b/stig/commands/tui/torrent.py @@ -197,3 +197,7 @@ class StopTorrentsCmd(base.StopTorrentsCmdbase, class VerifyTorrentsCmd(base.VerifyTorrentsCmdbase, mixin.polling_frenzy, mixin.make_request, mixin.select_torrents): provides = {'tui'} + + +class LabelCmd(base.LabelCmd, mixin.make_request, mixin.select_torrents, mixin.polling_frenzy): + provides = {'tui'} diff --git a/stig/completion/candidates.py b/stig/completion/candidates.py index bf04b178..1c310ede 100644 --- a/stig/completion/candidates.py +++ b/stig/completion/candidates.py @@ -371,6 +371,17 @@ async def objects_getter(**_): objects_getter=objects_getter, items_getter=None, filter_names=filter_names) +async def labels(curarg): + """All labels""" + labels = set() + if curarg != "": + response = await objects.srvapi.torrent.torrents(None, ('labels', ), from_cache=True) + [ + [labels.add(l) for l in t["labels"] if l.startswith(curarg)] + for t in response.torrents + ] + return Candidates(list(labels), label=curarg,curarg_seps=',') + async def _filter(curarg, filter_cls_name, objects_getter, items_getter, filter_names): """ diff --git a/stig/views/details.py b/stig/views/details.py index df9c2064..d6f96ee6 100644 --- a/stig/views/details.py +++ b/stig/views/details.py @@ -167,6 +167,10 @@ def __init__(self, label, needed_keys, human_readable=None, machine_readable=Non needed_keys=('name',)), Item('ID', needed_keys=('id',)), + Item('Labels', + needed_keys=('labels',), + human_readable=lambda t: ', '.join(t['labels']), + machine_readable=lambda t: ','.join(t['labels'])), Item('Hash', needed_keys=('hash',)), Item('Size', diff --git a/stig/views/torrent.py b/stig/views/torrent.py index a011db61..e1ed0f39 100644 --- a/stig/views/torrent.py +++ b/stig/views/torrent.py @@ -413,6 +413,19 @@ def get_value(self): COLUMNS['tracker'] = Tracker +class Labels(ColumnBase): + header = {'left': 'Labels'} + width = 10 + min_width = 5 + needed_keys = ('labels',) + align = 'left' + + def get_value(self): + ls = list(self.data['labels']) + ls.sort() + return ', '.join(ls) + +COLUMNS['labels'] = Labels class _TimeBase(ColumnBase): width = 10