#!/usr/bin/env python3 import argparse import os import pickle from tvsort import shows from tvsort.parser import EpisodeParseException, parse_episode, sub_bucket_name from appdirs import user_config_dir from fuzzywuzzy import fuzz from tabulate import tabulate from pprint import pprint from collections import namedtuple MatchedEpisode = namedtuple("MatchedEpisode", "root ep dest subdest score") """ Struct describing the intent to sort and episode file into a location root : abs path to the folder containing ep.file ep : associated EpInfo object dest : associated Show object score : scoring value Show::match returned """ def main(): " parse command line args " parser = argparse.ArgumentParser(description="sort tv shows") parser.add_argument("-s", "--src", nargs="+", help="", required=True) parser.add_argument("-d", "--dest", nargs="+", help="", required=True) parser.add_argument("--soft", action="store_true", help="Soft link instead of hard link") parser.add_argument("-r", "--rescan", action="store_true", help="Rescan library instead of using cache") parser.add_argument("--match-thresh", type=int, default=65) parser.add_argument("--mappings", nargs="+", default=[]) # many foo=bar transformations to help witch mapping args = parser.parse_args() if args.match_thresh <= 0 or args.match_thresh > 100: parser.error("--match-thresh must be 1-100") # mappings allow simple string transforms as workarounds for poorly named episodes mappings = {} for item in args.mappings: key, value = item.split("=") mappings[key] = value # load the library, an index of shows already sorted. the dirnames will be compared to incoming files cachedir = user_config_dir("tvsort") os.makedirs(cachedir, exist_ok=True) cache_file = os.path.join(cachedir, "library.cache") library = None if os.path.exists(cache_file) and not args.rescan: with open(cache_file, "rb") as f: try: library = pickle.load(f) except: print("Failed to load library cache") if not library: library = shows.create_index(args.dest) with open(cache_file, "wb") as f: pickle.dump(library, f) failures = [] results = [] # iterate through all children of the src dirs for srcdir in args.src: for fname in os.listdir(srcdir): # TODO season dirs are ignored for now if not os.path.isfile(os.path.join(srcdir, fname)): continue # Apply manually specified transformations item = fname for old, new in mappings.items(): item = item.replace(old, new) # Parse information from the episode file name try: epinfo, item = parse_episode(item) except EpisodeParseException: failures.append(fname) # Find a show from the library best matching this episode match_score = 0 best_match_show = None for show in library: value = fuzz.token_set_ratio(show.name.lower(), item.lower()) #TODO add algorithm swap arg for snakeoil if value > match_score: match_score = value best_match_show = show if match_score >= args.match_thresh: results.append( MatchedEpisode(srcdir, epinfo, best_match_show, sub_bucket_name(best_match_show, epinfo.major, epinfo.minor, epinfo.extra), match_score)) else: failures.append(fname) before = len(results) results = list( filter( lambda r: not os.path.exists(os.path.join(r.dest.root, r.dest.dir, r.subdest, r.ep.file)), results)) already_there = before - len(results) go = False while not go: tab_rows = [] i = 0 results.sort(key=lambda x: x.dest.dir) for item in results: row = [i, item.ep.file, item.ep.major, item.ep.minor, os.path.join(item.dest.dir, item.subdest) + "/", item.score, "soft" if args.soft else "hard"] tab_rows.append(row) i += 1 print(tabulate(tab_rows, headers=["number", "file", "season", "episode", "dest", "score", "link"])) if already_there: print("\n{} already in library and ignored".format(already_there)) if failures: print("\n") print("Could not match:") pprint(failures) if not results: print("no candidates for linking found!") return resp = input("create links? [y/N/]: ").lower().strip() if not resp or resp == "n": return if resp == "y": break exclude = [] for number in resp.split(): exclude.append(int(number)) exclude.sort(reverse=True) for number in exclude: results.pop(number) link = os.symlink if args.soft else os.link for item in results: src = os.path.join(item.root, item.ep.file) destdir = os.path.join(item.dest.root, item.dest.dir, item.subdest) dest = os.path.join(destdir, item.ep.file) # print("mkdir ", destdir) os.makedirs(destdir, exist_ok=True) # print(src, " -> ", dest) link(src, dest) if __name__ == '__main__': main()