163 lines
5.6 KiB
Python
163 lines
5.6 KiB
Python
#!/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/<lines to skip and print again>]: ").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()
|