add rank_df util

Author: chapmanjacobd

library/mediafiles/torrents_stop.py

@@ -3,7 +3,8 @@
 
 from library import usage
 from library.mediafiles.torrents_start import start_qBittorrent
-from library.utils import arggroups, consts, devices, iterables, nums, path_utils, printing
+from library.playback.torrents_info import qbt_get_tracker
+from library.utils import arggroups, consts, devices, nums, path_utils, printing
 from library.utils.log_utils import log
 
 
@@ -88,13 +89,8 @@ def torrents_stop():
                 new_path = Path(path_utils.mountpoint(torrent.content_path)) / new_path
 
             if args.tracker_dirnames:
-                tracker = torrent.tracker
-                if not tracker:
-                    tracker = iterables.safe_unpack(
-                        tr.url for tr in qbt_client.torrents_trackers(torrent.hash) if tr.url.startswith("http")
-                    )
-                if tracker:
-                    domain = path_utils.domain_from_url(tracker)
+                domain = qbt_get_tracker(qbt_client, torrent)
+                if domain:
                     new_path /= domain
 
             new_path.mkdir(parents=True, exist_ok=True)

library/mediafiles/torrents_stop_incomplete.py

@@ -3,7 +3,8 @@
 
 from library import usage
 from library.mediafiles.torrents_start import start_qBittorrent
-from library.utils import arggroups, consts, devices, iterables, path_utils, printing, strings
+from library.playback.torrents_info import qbt_get_tracker
+from library.utils import arggroups, consts, devices, path_utils, printing, strings
 from library.utils.log_utils import log
 
 
@@ -114,13 +115,8 @@ def torrents_stop_incomplete():
                 new_path = Path(path_utils.mountpoint(torrent.content_path)) / new_path
 
             if args.tracker_dirnames:
-                tracker = torrent.tracker
-                if not tracker:
-                    tracker = iterables.safe_unpack(
-                        tr.url for tr in qbt_client.torrents_trackers(torrent.hash) if tr.url.startswith("http")
-                    )
-                if tracker:
-                    domain = path_utils.domain_from_url(tracker)
+                domain = qbt_get_tracker(qbt_client, torrent)
+                if domain:
                     new_path /= domain
 
             new_path.mkdir(parents=True, exist_ok=True)

library/multidb/allocate_torrents.py

@@ -100,7 +100,7 @@ def allocate_torrents():
     torrents = list(args.db.query(*sqlgroups.playlists_fs_sql(args, limit=None)))
     total_size = sum(d["size"] for d in torrents)
     print(f"{len(torrents)} undownloaded torrents. {strings.file_size(total_size)} total space")
-    iterables.count_category(torrents, "tracker")
+    iterables.list_dict_value_counts(torrents, "tracker")
 
     if not torrents:
         processes.no_media_found()

library/playback/torrents_info.py

@@ -19,6 +19,15 @@ def parse_args():
     return args
 
 
+def qbt_get_tracker(qbt_client, torrent):
+    tracker = torrent.tracker
+    if not tracker:
+        tracker = iterables.safe_unpack(
+            tr.url for tr in qbt_client.torrents_trackers(torrent.hash) if tr.url.startswith("http")
+        )
+    return domain_from_url(tracker)
+
+
 def torrents_info():
     args = parse_args()
 
@@ -70,12 +79,7 @@ def torrents_info():
 
     torrents_by_tracker = {}
     for torrent in all_torrents:
-        tracker = torrent.tracker
-        if not tracker:
-            tracker = iterables.safe_unpack(
-                tr.url for tr in qbt_client.torrents_trackers(torrent.hash) if tr.url.startswith("http")
-            )
-        torrents_by_tracker.setdefault(domain_from_url(tracker), []).append(torrent)
+        torrents_by_tracker.setdefault(qbt_get_tracker(qbt_client, torrent), []).append(torrent)
 
     interesting_states = [
         "stoppedUP",

library/utils/iterables.py

@@ -1,4 +1,5 @@
 import math
+from collections import Counter
 from collections.abc import Iterable, Iterator
 from functools import wraps
 from typing import Any
@@ -149,7 +150,7 @@ def list_dict_unique(data: list[dict], unique_keys: list[str]) -> list[dict]:
     return list_
 
 
-def count_category(list_of_dicts, key_name):
+def list_dict_value_counts(list_of_dicts, key_name):
     category_counts = {}
     for item in list_of_dicts:
         category = item.get(key_name)
@@ -251,3 +252,8 @@ def zipkw(**kwargs):
 
     for combination in zip(*values):
         yield dict(zip(keys, combination))
+
+
+def value_counts(input_list):
+    counts = Counter(input_list)
+    return [counts[item] for item in input_list]

library/utils/pd_utils.py

@@ -49,3 +49,39 @@ def available_name(df, column_name):
         else:
             column_name = f"{column_name}_1"
     return column_name
+
+
+def rank_dataframe(df, column_weights):
+    """
+    ranked_df = rank_dataframe(
+        df,
+        column_weights={
+            "progress": {"direction": "desc", "weight": 6},
+            "size": {"direction": "asc", "weight": 3}
+        }
+    )
+    """
+    ranks = df[column_weights.keys()].apply(
+        lambda x: x.rank(
+            method="min",
+            na_option="bottom",
+            ascending=column_weights.get(x.name, {}).get("direction") == "asc",
+        )
+        * column_weights.get(x.name, {}).get("weight", 1),
+    )
+
+    unranked_columns = set(df.select_dtypes(include=["number"]).columns) - set(ranks.columns)
+    if unranked_columns:
+        print(
+            "Unranked columns:\n"
+            + "\n".join([f"""    "{s}": {{ 'direction': 'desc' }}, """ for s in unranked_columns]),
+        )
+
+    scaled_ranks = (ranks - 1) / (len(ranks.columns) - 1)
+    scaled_df = df.iloc[scaled_ranks.sum(axis=1).sort_values().index]
+    return scaled_df.reset_index(drop=True)
+
+
+def count_category(df, key_name):
+    df[f"{key_name}_count"] = df.groupby(key_name)[key_name].transform("size")
+    return df