#!/usr/bin/python3 # """Tool to query TheMovieDB and rename files.""" # # Copyright (C) 2019 Nicolas Schodet # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to # deal in the Software without restriction, including without limitation the # rights to use, copy, modify, merge, publish, distribute, sublicense, and/or # sell copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS # IN THE SOFTWARE. # import argparse import configparser import os import re import subprocess import requests import tmdbsimple as tmdb CONFIG = "~/.tmdblookup.conf" def title_year(title, year): """Return title, with year if available.""" if year is not None: return f"{title} ({year})" else: return title class Movie: """A movie.""" def __init__(self, title, year, i): self.title = title self.year = year self.i = i def __str__(self): return title_year(self.title, self.year) class Episode: """A TV show episode.""" def __init__(self, season, episode, title): self.season = season self.episode = episode self.title = title def __str__(self): return f"{self.season}x{self.episode:02d} {self.title}" class TVShow: """A TV show with its episodes.""" def __init__(self, title=None, year=None, i=None): self.title = title self.year = year self.i = i self.episodes = [] def load(self, episodes_file): """Load a TV show and its episodes from a file.""" title = episodes_file.readline().strip() if not title: raise RuntimeError("no TV show title in file") year = None m = re.fullmatch(r"(.*) \(([0-9]+)\)", title) if m: title = m.group(1) year = int(m.group(2)) episodes = [] for line in episodes_file: m = re.fullmatch(r"([0-9]+)x([0-9]+) (.*)", line.strip()) if not m: raise RuntimeError("wrong episodes file format") episodes.append(Episode(int(m.group(1)), int(m.group(2)), m.group(3))) if not episodes: raise RuntimeError("no episode in file") self.title = title self.year = year self.episodes = episodes def filter(self, seasons, episodes): """Filter a list of episodes according to parameters.""" filtered = [] for e in self.episodes: if seasons is not None and e.season not in seasons: continue if episodes is not None and (e.season, e.episode) not in episodes: continue filtered.append(e) self.episodes = filtered def __str__(self): r = [title_year(self.title, self.year)] for e in self.episodes: r.append(str(e)) return "\n".join(r) def parse(args): """Parse arguments and configuration.""" defaults = {} # Parse configuration. config = configparser.ConfigParser() config.read([os.path.expanduser(CONFIG)]) defaults.update(dict(config.items("DEFAULT"))) # Parse options. p = argparse.ArgumentParser(description=__doc__) p.set_defaults(**defaults) p.add_argument("--api-key", help="API key created in your account") p.add_argument("-l", "--language", help="content language") p.add_argument( "-t", "--tv-show", action="store_true", help="query TV shows instead of movies", ) p.add_argument( "-s", "--seasons", nargs="+", type=int, metavar="S", help="select TV show seasons", ) p.add_argument( "-a", "--all-seasons", action="store_true", help="select all TV show seasons", ) p.add_argument( "-e", "--episodes", nargs="+", metavar="SxE", help="select TV show episodes, use SxE notation (e.g. 1x2 for the second" " episode of the first season), or just the episode number if season is" " selected with -s", ) p.add_argument( "-f", "--file", type=argparse.FileType("rt"), help="load TV show episodes from given file instead of querying the" " database (same format as this tool output)", ) p.add_argument( "-i", "--interactive", action="store_true", help="ask the user when not sure (else abort)", ) p.add_argument( "-r", "--rename", nargs="+", metavar="FILE", help="rename files given as parameter", ) p.add_argument( "-n", "--dry-run", action="store_true", help="only pretend, but do nothing" ) p.add_argument("--suffix", help="suffix to add to each file being renamed") p.add_argument("--year", help="query year") p.add_argument("query", nargs="*", help="query terms") options = p.parse_args() # Check and return options. if options.api_key is None and options.file is None: p.error("need an API key") if not options.query and options.file is None: p.error("need a query or an episodes file") if (options.query or options.year is not None) and options.file is not None: p.error("can not give both a query and an episodes file") options.query = " ".join(options.query) if options.year is None: m = re.fullmatch(r"(.*) \(([0-9]+)\)", options.query) if m: options.query = m.group(1) options.year = m.group(2) if options.episodes is not None: episodes = [] for e in options.episodes: m = re.fullmatch(r"([0-9]+)x([0-9]+)", e) if m: s, e = int(m.group(1)), int(m.group(2)) else: try: e = int(e) except ValueError: raise RuntimeError("bad episode number format") if options.seasons is None or len(options.seasons) != 1: raise RuntimeError( "need a single season, or full episode number (SxE format)" ) s = options.seasons[0] episodes.append((s, e)) options.episodes = episodes if ( options.file is not None or options.seasons is not None or options.all_seasons or options.episodes is not None ): # Imply TV show if episodes are selected. options.tv_show = True elif options.rename and options.tv_show: # When renaming TV shows, if no episodes selection, select all episodes. options.all_seasons = True return options def search(query, options): """Perform a TV show or movie search.""" tmdb.API_KEY = options.api_key search = tmdb.Search() searchd = dict(query=query) if options.language: searchd["language"] = options.language results = [] if not options.tv_show: if options.year is not None: searchd["year"] = options.year search.movie(**searchd) for r in search.results: y = None title = r["title"] try: y = r["release_date"] except KeyError: pass if y: y, _, _ = y.split("-") i = r["id"] results.append(Movie(title, y, i)) else: if options.year is not None: searchd["first_air_date_year"] = options.year search.tv(**searchd) for r in search.results: y = None title = r["name"] try: y = r["first_air_date"] except KeyError: pass if y: y, _, _ = y.split("-") i = r["id"] results.append(TVShow(title, y, i)) return results def select(results, header=None): """Select one of the results.""" if header: header = [f"--header={header}"] height = min(12, len(results) + 3) else: header = [] height = min(12, len(results) + 2) r = subprocess.run( ["fzf", f"--height={height}", "--with-nth=2", "--delimiter=\t"] + header, input="\n".join(f"{i}\t{r}" for i, r in enumerate(results)), stdout=subprocess.PIPE, text=True, ) if r.returncode != 0: return None else: i = int(r.stdout.split(maxsplit=1)[0]) return results[i] def get_unique(query, options): """Return a TV show or movie, ask the user if needed.""" results = search(query, options) if len(results) == 0: raise RuntimeError("no result") elif len(results) == 1: return results[0] else: r = select(results) if r is None: raise RuntimeError("no selection done") return r def get_seasons(tv_show, options): """Get list of seasons for a TV show.""" if options.episodes is not None: seasons = set() for s, e in options.episodes: seasons.add(s) return list(seasons) elif options.seasons is not None: return options.seasons else: tv = tmdb.TV(tv_show.i) infod = dict() if options.language: infod["language"] = options.language tv.info(**infod) return [s["season_number"] for s in tv.seasons] def add_episodes(tv_show, options): """Populate a TV show with its episodes.""" for s in get_seasons(tv_show, options): season = tmdb.TV_Seasons(tv_show.i, s) infod = dict() if options.language: infod["language"] = options.language try: season.info(**infod) except requests.exceptions.HTTPError: raise RuntimeError("season %d does not exists" % s) for e in season.episodes: tv_show.episodes.append(Episode(s, e["episode_number"], e["name"])) tv_show.filter(options.seasons, options.episodes) def rename_episode(f, tv_show, episode, options): """Rename a single episode file.""" directory = title_year(tv_show.title, tv_show.year) _, ext = os.path.splitext(f) suffix = f" - {options.suffix}" if options.suffix else "" name = f"{tv_show.title} {episode}{suffix}{ext}" newf = os.path.join(directory, name) print(f"rename {repr(f)} to {repr(newf)}") if not options.dry_run: os.renames(f, newf) def rename_tv_shows(results, options): """Rename TV shows files.""" assert len(results) == 1 tv_show = results[0] if options.interactive: for f in options.rename: e = select(tv_show.episodes, f"select episode for {repr(f)}") if e is None: print(f"skipped {f}") else: rename_episode(f, tv_show, e, options) else: n_files = len(options.rename) n_episodes = len(tv_show.episodes) if n_episodes != n_files: raise RuntimeError( f"number of files ({n_files}) and number of episodes ({n_episodes}) do not match" ) for f, e in zip(options.rename, tv_show.episodes): rename_episode(f, tv_show, e, options) def rename_movie(f, movie, options): """Rename a movie file.""" _, ext = os.path.splitext(f) suffix = f" - {options.suffix}" if options.suffix else "" newf = f"{movie}{suffix}{ext}" print(f"rename {repr(f)} to {repr(newf)}") if not options.dry_run: os.renames(f, newf) def rename_movies(results, options): """Rename movies files.""" if options.interactive or len(options.rename) != 1 or len(results) != 1: for f in options.rename: movie = select(results, f"select movie for {repr(f)}") if movie is None: print(f"skipped {f}") else: rename_movie(f, movie, options) else: rename_movie(options.rename[0], results[0], options) def run(options): """Run requested action.""" # Find movies, TV shows or episodes. if options.file is not None: r = TVShow() r.load(options.file) r.filter(options.seasons, options.episodes) results = [r] elif ( options.seasons is not None or options.all_seasons or options.episodes is not None ): r = get_unique(options.query, options) add_episodes(r, options) results = [r] else: results = search(options.query, options) # Action: print or rename. if not options.rename: for r in results: print(r) elif options.tv_show: rename_tv_shows(results, options) else: rename_movies(results, options) if __name__ == "__main__": import sys try: options = parse(sys.argv[1:]) run(options) except RuntimeError as e: sys.exit(str(e)) except OSError as e: sys.exit(str(e))