telegram-video-downloader/src/vid_utils.py

import os
import logging
from glob import glob, escape
from subprocess import Popen, PIPE
from contextlib import contextmanager

from telegram import InlineKeyboardButton

from db import VidDatabase

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)

class BadLink(Exception):
    pass


class Video:
    def __init__(self, link=None, vid=None, init_keyboard=False):
        self.db = VidDatabase(os.path.join(os.environ['CONF_FOLDER'], "viddb.sqlite3"))
        if not self.db.is_valid:
            # Database file not present
            # Create a new database
            self.db.create()

        if vid is None and link is not None:
            self.link = link
            self.file_name = None
        elif vid is not None and link is None:
            self.link, self.code = self.db.select_vid(vid)
        else:
            raise Exception('what is going on?')

        if init_keyboard:
            self.formats = self.get_formats()
            self.keyboard = self.generate_keyboard()

    def get_formats(self):
        formats = {}

        p = Popen(["youtube-dl", "-F", self.link], stdout=PIPE, stderr=PIPE).communicate()
        it = iter(str(p[0], 'utf-8').split('\n'))  # iterator of output lines

        try:
            while "code  extension" not in next(it):
                pass  # Remove garbage lines
        except StopIteration:
            raise BadLink("youtube-dl couldn't download the link you provided")  # Isn't a valid youtube link

        while True:
            try:
                line = next(it)
                if not line:
                    raise StopIteration  # Usually the last line is empty
                if "video only" in line:
                    continue  # I don't need video without audio
            except StopIteration:
                break
            else:
                format_code, extension, resolution, *_ = line.strip().split()
                key = '{},{}'.format(extension, resolution)
                index = self.db.insert_vid(self.link, format_code)
                formats[key] = index
        logger.info("Fromats: {}".format(formats))
        return formats

    def generate_keyboard(self):
        """ Generate a list of InlineKeyboardButton of resolutions """
        kb = []

        for key in self.formats.keys():
            cb = "{}".format(self.formats[key])
            kb.append([InlineKeyboardButton(key, callback_data=cb)])
        return kb

    def download(self):
        logger.info("Downloading {}".format(self.link))
        p = Popen(["youtube-dl", "-o", "/bot/out/%(title)s-%(id)s.%(ext)s", "-f", self.code, self.link], stdout=PIPE, stderr=PIPE).communicate()

        for line in str(p[0], 'utf-8').split('\n'):
            logger.info(line)
            if "[download] Destination:" in line:
                self.file_name = line[24:]  # name of the file

    def check_dimension(self):
        try:
            if os.path.getsize(self.file_name) > 50 * 1024 * 1023:
                Popen(["split", "-b", "49M", self.file_name, self.file_name])
                os.remove(self.file_name)
            return glob(escape(self.file_name) + '*')
        except AttributeError as e:
            logger.error(e)
            return None

    @contextmanager
    def send(self):
        files = self.check_dimension()  # split if size >= 50MB
        yield files
        for f in files:  # removing old files
            os.remove(f)





#__________________________OLD STUFFS, TOUCH CAREFULLY__________________________

# this is the soft-split version, require avconv, but the audio isn't synchronized, avconv's problems :(
'''
def get_duration(filepath): # get duration in seconds
    cmd = "avconv -i %s" % filepath
    p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE)
    di = p.communicate()
    for line in di:
        if line.rfind(b"Duration") > 0:
            duration = str(re.findall(b"Duration: (\d+:\d+:[\d.]+)", line)[0])
    return 3600 * int(duration[2: 4]) + 60 * int(duration[5: 7]) + int(duration[8: 10])

def check_dimension(f): # if f is bigger than 50MB split it in subvideos
    if os.path.getsize(f) > 50 * 1024 * 1023:
        duration = get_duration(f)
        for i in range(0, duration, 180):
            start = strftime("%H:%M:%S", strptime('{0} {1} {2}'.format(i // 3600, (i // 60) % 60, i % 60), "%H %M %S")) # TODO this is  not pythonic code!
            os.system("""avconv -i '{0}' -vcodec copy -acodec copy -ss {1} -t {2} 'part_{3}.mp4'""".format(f, start, 180, (i // 180) % 180))
        os.remove(f) # delete original file
'''