Use custom ffmpeg class (#152)

* Add ffmpeg-dl

Heavily borrowed from https://stackoverflow.com/a/77052858/3939155. Will fix #151 as we avoid socket.AF_UNIX and just read stdout

* Address TODOs

* Perform final refresh and close

* Reorder imports

Author: AlexAplin

Diff for: nndownload/ffmpeg_dl.py

@@ -0,0 +1,81 @@
+import regex as re
+import subprocess
+import warnings
+from datetime import timedelta, datetime
+from typing import AnyStr, List
+
+import ffmpeg
+from tqdm import TqdmExperimentalWarning
+from tqdm.rich import tqdm_rich
+
+warnings.filterwarnings("ignore", category=TqdmExperimentalWarning)
+
+
+class FfmpegDLException(Exception):
+    """Raised when a download fails."""
+    pass
+
+
+class FfmpegDL:
+    """Send input streams for download to an `ffmpeg` subprocess."""
+
+    FF_GLOBAL_ARGS = [
+        "-progress",
+        "-",
+        "-nostats",
+        "-y"
+    ]
+
+    REGEX_TIME_GROUP = "([0-9]{2}:[0-9]{2}:[0-9]{2}[.[0-9]*]?)"
+    REGEX_OUT_TIME = re.compile(
+        r"out_time=[ ]*" + REGEX_TIME_GROUP
+    )
+
+    @classmethod
+    def get_timedelta(cls, time_str: AnyStr, str_format: AnyStr = "%H:%M:%S.%f"):
+        t = datetime.strptime(time_str, str_format)
+        return timedelta(hours=t.hour, minutes=t.minute, seconds=t.second, microseconds=t.microsecond)
+
+    def __init__(self, streams: List, input_kwargs: List, output_path: AnyStr, output_kwargs: List, global_args: List = FF_GLOBAL_ARGS):
+        inputs = []
+        for stream in streams:
+            input = ffmpeg.input(stream, **input_kwargs)
+            inputs.append(input)
+        stream_spec = ffmpeg.output(*inputs, output_path, **output_kwargs).global_args(*global_args)
+
+        self.proc_args = ffmpeg._run.compile(stream_spec=stream_spec)
+        self.proc: subprocess.Popen = None
+
+    def load_subprocess(self):
+        self.proc = subprocess.Popen(
+            args=self.proc_args,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=False,
+        )
+
+    def convert(self, name: AnyStr, duration: float):
+        progress = tqdm_rich(desc=name, unit="seg", colour="green", total=duration)
+
+        self.load_subprocess()
+
+        stdout_line = None
+        while True:
+            if self.proc.stdout is None:
+                continue
+            prev_line = stdout_line
+            stdout_line = self.proc.stdout.readline().decode("utf-8", errors="replace").strip()
+            out_time_data = self.REGEX_OUT_TIME.search(stdout_line)
+            if out_time_data is not None:
+                out_time = self.get_timedelta(out_time_data.group(1))
+                progress.update(out_time.total_seconds() - progress.n)
+                continue
+            if stdout_line == "" and self.proc.poll() is not None:
+                progress.refresh()
+                progress.close()
+                exit_code = self.proc.poll()
+                if exit_code:
+                    raise FfmpegDLException(prev_line)
+                else:
+                    break

Diff for: nndownload/nndownload.py

@@ -14,28 +14,23 @@
 import os
 import re
 import shutil
-import socket
 import sys
 import tempfile
 import threading
 import time
-import warnings
 import xml.dom.minidom
 from typing import AnyStr, List, Match
 
 import aiohttp
 import requests
-import ffmpeg
-import gevent.monkey
-gevent.monkey.patch_all(ssl=False,thread=False)
 from aiohttp_socks import ProxyConnector
 from bs4 import BeautifulSoup
 from mutagen.mp4 import MP4, MP4StreamInfoError
 from requests.adapters import HTTPAdapter
 from requests.utils import add_dict_to_cookiejar
 from urllib3.util import Retry
-from tqdm import tqdm, TqdmWarning
 
+from ffmpeg_dl import FfmpegDL, FfmpegDLException
 
 __version__ = "1.16"
 __author__ = "Alex Aplin"
@@ -1250,80 +1245,24 @@ def download_video_part(session: requests.Session, start, end, filename: AnyStr,
             update_multithread_progress(len(block))
 
 
-def capture_ffmpeg_progress(filename, sock, handler):
-    """Capture and send ffmpeg events to the provided handler."""
-
-    connection, client_address = sock.accept()
-    data = b""
-    try:
-        while True:
-            more_data = connection.recv(16)
-            if not more_data:
-                break
-            data += more_data
-            lines = data.split(b"\n")
-            for line in lines[:-1]:
-                line = line.decode()
-                parts = line.split("=")
-                key = parts[0] if len(parts) > 0 else None
-                value = parts[1] if len(parts) > 1 else None
-                handler(key, value)
-            data = lines[-1]
-    finally:
-        connection.close()
-
-
-@contextlib.contextmanager
-def watch_ffmpeg_progress(handler):
-    """Spawn a socket to capture ffmpeg events."""
-
-    with get_temp_dir() as temp_dir:
-        socket_filename = os.path.join(temp_dir, "sock")
-        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        with contextlib.closing(sock):
-            sock.bind(socket_filename)
-            sock.listen(1)
-            child = gevent.spawn(capture_ffmpeg_progress, socket_filename, sock, handler)
-            try:
-                yield socket_filename
-            except:
-                gevent.kill(child)
-                raise
-
-
-@contextlib.contextmanager
-def show_ffmpeg_progress(duration: float):
-    """Render a tqdm progress bar relative to a specified stream duration."""
-
-    warnings.filterwarnings("ignore", category=TqdmWarning)
-    with tqdm(total=round(duration, 2)) as progress_bar:
-        def handler(key, value):
-            if key == "out_time_ms":
-                time = round(float(value) / 1000000., 2)
-                progress_bar.update(time - progress_bar.n)
-            elif key == "progress" and value == "end":
-                progress_bar.update(progress_bar.total - progress_bar.n)
-        with watch_ffmpeg_progress(handler) as socket_filename:
-            yield socket_filename
-
-
-def perform_ffmpeg_dl(filename: AnyStr, duration: float, streams: List):
+def perform_ffmpeg_dl(video_id: AnyStr, filename: AnyStr, duration: float, streams: List):
     """Send video and/or audio stream to ffmpeg for download."""
 
-    inputs = []
     try:
-        with show_ffmpeg_progress(duration) as socket_filename:
-            for stream in streams:
-                input = ffmpeg.input(stream, protocol_whitelist="https,http,tls,tcp,file,crypto", allowed_extensions="ALL")
-                inputs.append(input)
-            output = ffmpeg.output(*inputs, filename, vcodec="copy", acodec="copy")
-            output = output.global_args("-progress", "unix://{}".format(socket_filename), "-y")
-            output.run(capture_stdout=True, capture_stderr=True)
-            return True
-    except ffmpeg.Error as error:
-        stderr = error.stderr.decode("utf8")
-        actual_error = stderr.splitlines()[-1]
-        raise FormatNotAvailableException(f"ffmpeg exited with an error: {actual_error}")
+        video_download = FfmpegDL(streams=streams,
+                                    input_kwargs={
+                                        "protocol_whitelist": "https,http,tls,tcp,file,crypto",
+                                        "allowed_extensions": "ALL",
+                                    },
+                                    output_path=filename,
+                                    output_kwargs={
+                                        "vcodec": "copy",
+                                        "acodec": "copy",
+                                    })
+        video_download.convert(name=video_id, duration=duration)
+        return True
+    except FfmpegDLException as error:
+        raise FormatNotAvailableException(f"ffmpeg failed to download the video or audio stream with the following error: \"{error}\"")
     except Exception:
         raise FormatNotAvailableException("Failed to download video or audio stream")
 
@@ -1366,7 +1305,7 @@ def download_video_media(session: requests.Session, filename: AnyStr, template_p
                     generic_dl_request(session, key_url, key_path, binary=True)
                     rewrite_file(m3u8_path, key_url, key_path)
                     m3u8_streams.append(m3u8_path)
-            continue_code = perform_ffmpeg_dl(filename, float(template_params["duration"]), m3u8_streams)
+            continue_code = perform_ffmpeg_dl(template_params["id"], filename, float(template_params["duration"]), m3u8_streams)
             os.rename(filename, complete_filename)
             return continue_code