Interface for transmission's labels feature

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)

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',

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):

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),

stig/commands/base/torrent.py

@@ -32,7 +32,8 @@ class AddTorrentsCmdbase(metaclass=CommandMeta):
     description = 'Download torrents'
     usage = ('add [<OPTIONS>] <TORRENT> <TORRENT> <TORRENT> ...',)
     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 [<OPTIONS>] <TORRENT FILTER> <TORRENT FILTER>... <[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)

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'}

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

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'}

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):
     """

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',