mediathek_download/main.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import os
import re
import shutil
import tempfile
# import urllib.request
from configparser import ConfigParser

import feedparser
import requests
import tvdbsimple as tvdb
from fuzzywuzzy import fuzz
from pymkv import MKVFile, MKVTrack

logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.DEBUG)

seasons = ConfigParser(interpolation=None)
seasons.read('seasons.ini')
config = ConfigParser()
config.read('config.ini')

tvdb.KEYS.API_KEY = config['TheTVDB']['apikey']


def get_match(query,
              choices,
              key=lambda x: x,
              score_cutoff=0,
              scorer=fuzz.ratio):
    res = [(item, scorer(query, key(item))) for item in choices]
    return sorted([x for x in res if x[1] >= score_cutoff],
                  key=lambda x: x[1],
                  reverse=True)


def get_episod_info(season_id, name):
    episodes = tvdb.Series_Episodes(season_id, 'de').all()
    matches = get_match(name,
                        episodes,
                        key=lambda x: x['episodeName'],
                        score_cutoff=90)

    if len(matches) >= 1:
        titel_dict = {}
        titel_dict['title'] = matches[0][0]['episodeName']
        titel_dict['season'] = str(matches[0][0]['airedSeason'])
        titel_dict['episod'] = str(matches[0][0]['airedEpisodeNumber'])
        return titel_dict


def build_filename(name: str, titel_dict: dict):
    if not titel_dict:
        return None
    if titel_dict['episod'].isdigit():
        filename = '{name} S{season:>02}E{episod:>02} {title}'.format(
            name=name, **titel_dict)
    else:
        filename = '{name} S{season:>02}.{episod} {title}'.format(name=name,
                                                                  **titel_dict)

    return filename.strip().replace(' ', '.')


def make_mkv(data):
    series_name = data['series']
    default_lang = seasons[series_name]['default_language']
    lang_sort = seasons[series_name]['languages'].split(',')

    def add_video(lang, path):
        logging.debug('[MKV] add video (%s, %s)', lang, path)
        video = MKVTrack(path, track_id=0, language=lang, default_track=True)
        mkv.add_track(video)

    def add_audio(lang, path):
        logging.debug('[MKV] add audio (%s, %s)', lang, path)
        audio = MKVTrack(path,
                         track_id=1,
                         language=lang,
                         default_track=lang == default_lang)
        mkv.add_track(audio)

    mkv = MKVFile()
    add_video(default_lang, data['lang'][default_lang])

    for lang in lang_sort:
        mp4_file = data['lang'].pop(lang, None)
        if mp4_file is not None:
            add_audio(lang, mp4_file)

    mkv.mux(data['mkvpath'])


def parse_feed(series_name):
    d = feedparser.parse(seasons[series_name]['rss'])

    episodes = dict()
    for item in d['entries']:
        lang = 'ger'  # init lang
        title = item['title']
        link = item['link']
        for lang in seasons[series_name]['languages'].split(','):
            pattern = seasons[series_name][lang]
            season_id = seasons.get(series_name, 'thetvdb', fallback=None)
            groups = [x.groupdict() for x in re.finditer(pattern, title, re.M)]
            if len(groups) == 1:
                break

        if groups == []:
            logging.warning('skip %s', title)
            continue

        titel_dict = groups[0]
        if season_id is not None:
            titel_dict = get_episod_info(season_id, titel_dict['title'])
        if not titel_dict:
            logging.warning('skip %s', title)
            continue
        filename = build_filename(series_name, titel_dict)
        if filename not in episodes:
            episodes[filename] = dict()
            episodes[filename]['lang'] = dict()
        episodes[filename]['lang'][lang] = link
        season = '{season:>02}'.format(**titel_dict)
        episod = '{episod:>02}'.format(**titel_dict)
        basepath = os.path.join(seasons[series_name]['save_path'], series_name,
                                f'Season {season:>02}')
        mkvpath = os.path.join(basepath, filename + '.mkv')

        episodes[filename]['season'] = season
        episodes[filename]['episod'] = episod
        episodes[filename]['series'] = series_name
        episodes[filename]['basepath'] = basepath
        episodes[filename]['mkvpath'] = mkvpath

    return episodes


def download_files(episodes, dryrun=False):
    for episod in sorted(episodes.keys()):
        item = episodes[episod]
        if os.path.exists(item['mkvpath']):
            continue

        basepath = item['basepath']
        os.makedirs(basepath, exist_ok=True)
        temp = tempfile.TemporaryDirectory(dir=basepath)

        logging.info('Episod: %s', episod)

        for lang in item['lang']:
            link = item['lang'][lang]
            filepath = os.path.join(temp.name, os.path.basename(link))

            logging.info('start downloading...')
            logging.debug('source: %s', link)
            logging.debug('destination: %s', filepath)
            if dryrun:
                continue
            try:
                # response = urllib.request.urlopen(link)
                with requests.get(link, stream=True) as r:
                    with open(filepath, 'wb') as f:
                        shutil.copyfileobj(r.raw, f)

                item['lang'][lang] = filepath
            except:
                logging.error("could not download file.")

        if dryrun:
            continue
        try:
            make_mkv(item)
        except:
            logging.error("could not build mkv.")
        temp.cleanup()


def run():
    logging.info('====== START ======')
    for series_name in seasons.sections():
        episodes = parse_feed(series_name)
        download_files(episodes, dryrun=False)

    logging.info('====== END ======')


if __name__ == "__main__":
    run()