better format detection

Dockerfile

@@ -1,6 +1,8 @@
 FROM python:3.8-slim-buster
 WORKDIR /bot
 COPY requirements.txt .
+RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg; \
+	rm -rf /var/lib/apt/lists/*;
+RUN pip3 install -U youtube-dl
 RUN pip3 install -r requirements.txt
-RUN pip3 install --upgrade youtube-dl
 COPY src src

src/main.py

@@ -17,6 +17,8 @@ def get_format(update, context):
         video = Video(link=update.message.text, init_keyboard=True)
     except BadLink as e:
         update.message.reply_text("Bad link: {}".format(e))
+    except Exception as e:
+        update.message.reply_text("Error: {}".format(e))
     else:
         reply_markup = InlineKeyboardMarkup(video.keyboard)
         update.message.reply_text('Choose format:', reply_markup=reply_markup)

src/vid_utils.py

@@ -1,5 +1,7 @@
 import os
 import logging
+from json import loads
+from json.decoder import JSONDecodeError
 from glob import glob, escape
 from subprocess import Popen, PIPE
 from contextlib import contextmanager
@@ -17,7 +19,7 @@ class BadLink(Exception):
 
 class Video:
     def __init__(self, link=None, vid=None, init_keyboard=False):
-        self.db = VidDatabase(os.path.join(os.environ['CONF_FOLDER'], "viddb.sqlite3"))
+        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
@@ -28,64 +30,76 @@ class Video:
             self.file_name = None
         elif vid is not None and link is None:
             self.link, self.code = self.db.select_vid(vid)
+            self.code, self.audio_only = self.code.split('|')
         else:
             raise Exception('what is going on?')
 
         if init_keyboard:
-            self.formats = self.get_formats()
-            self.keyboard = self.generate_keyboard()
+            try:
+                self.formats = self.get_formats()
+                self.keyboard = self.generate_keyboard()
+            except Exception:
+                raise
 
     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
+        p = Popen(['youtube-dl', '-J', self.link], stdout=PIPE, stderr=PIPE).communicate()
+        if b'ERROR' in p[1]:
+            raise Exception('video URL not supported')
 
-        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
+        video_info = loads(str(p[0], 'utf-8'))
 
-        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))
+        if video_info.get('_type', None) == 'playlist':
+            video_info = video_info.get('entries')[0]
+
+        if video_info.get('formats') is not None:
+            for vid in video_info.get('formats'):
+                self.add_format(formats, vid)
+        else:
+            self.add_format(formats, video_info)
+
+        logger.info('Formats: {}'.format(formats))
         return formats
 
+    def add_format(self, formats, vid):
+        format_code = vid.get('format_id')
+        extension = vid.get('ext')
+        name = vid.get('format')
+        key = '{},{}'.format(extension, name)
+        code = '{}|{}'.format(format_code, 'audio only' in name)
+        index = self.db.insert_vid(self.link, code)
+        formats[key] = index
+
     def generate_keyboard(self):
-        """ Generate a list of InlineKeyboardButton of resolutions """
+        ''' Generate a list of InlineKeyboardButton of resolutions '''
         kb = []
 
         for key in self.formats.keys():
-            cb = "{}".format(self.formats[key])
+            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()
+        logger.info('Downloading {}'.format(self.link))
+        cmd = ['youtube-dl', '-o', '/bot/out/%(title)s-%(id)s.%(ext)s']
+        if self.audio_only in [False, "False"]:
+            self.code = self.code + '+bestaudio'
+        cmd.extend(['-f', self.code, self.link])
+        logger.info(cmd)
+        p = Popen(cmd, stdout=PIPE, stderr=PIPE).communicate()
 
         for line in str(p[0], 'utf-8').split('\n'):
             logger.info(line)
-            if "[download] Destination:" in line:
+            if '[download] Destination:' in line:
                 self.file_name = line[24:]  # name of the file
+            if '[ffmpeg] Merging formats into' in line:
+                self.file_name = line[31:-1]  # 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])
+                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:
@@ -98,29 +112,3 @@ class Video:
         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
-'''