[websockets] Add WebSocketFragmentFD (#399)

Necessary for #392

Co-authored by: nao20010128nao, pukkandan

yt_dlp/YoutubeDL.py

@@ -127,13 +127,14 @@ from .downloader import (
 )
 from .downloader.rtmp import rtmpdump_version
 from .postprocessor import (
+    get_postprocessor,
+    FFmpegFixupDurationPP,
     FFmpegFixupM3u8PP,
     FFmpegFixupM4aPP,
     FFmpegFixupStretchedPP,
+    FFmpegFixupTimestampPP,
     FFmpegMergerPP,
     FFmpegPostProcessor,
-    # FFmpegSubtitlesConvertorPP,
-    get_postprocessor,
     MoveFilesAfterDownloadPP,
 )
 from .version import __version__
@@ -2723,6 +2724,8 @@ class YoutubeDL(object):
                     downloader = (get_suitable_downloader(info_dict, self.params).__name__
                                   if 'protocol' in info_dict else None)
                     ffmpeg_fixup(downloader == 'HlsFD', 'malformed AAC bitstream detected', FFmpegFixupM3u8PP)
+                    ffmpeg_fixup(downloader == 'WebSocketFragmentFD', 'malformed timestamps detected', FFmpegFixupTimestampPP)
+                    ffmpeg_fixup(downloader == 'WebSocketFragmentFD', 'malformed duration detected', FFmpegFixupDurationPP)
 
                 fixup()
                 try:

yt_dlp/compat.py

@@ -3030,6 +3030,21 @@ except AttributeError:
     compat_Match = type(re.compile('').match(''))
 
 
+import asyncio
+try:
+    compat_asyncio_run = asyncio.run
+except AttributeError:
+    def compat_asyncio_run(coro):
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        loop.run_until_complete(coro)
+
+    asyncio.run = compat_asyncio_run
+
+
 __all__ = [
     'compat_HTMLParseError',
     'compat_HTMLParser',
@@ -3037,6 +3052,7 @@ __all__ = [
     'compat_Match',
     'compat_Pattern',
     'compat_Struct',
+    'compat_asyncio_run',
     'compat_b64decode',
     'compat_basestring',
     'compat_chr',

yt_dlp/downloader/__init__.py

@@ -24,6 +24,7 @@ from .rtsp import RtspFD
 from .ism import IsmFD
 from .mhtml import MhtmlFD
 from .niconico import NiconicoDmcFD
+from .websocket import WebSocketFragmentFD
 from .youtube_live_chat import YoutubeLiveChatReplayFD
 from .external import (
     get_external_downloader,
@@ -42,6 +43,7 @@ PROTOCOL_MAP = {
     'ism': IsmFD,
     'mhtml': MhtmlFD,
     'niconico_dmc': NiconicoDmcFD,
+    'websocket_frag': WebSocketFragmentFD,
     'youtube_live_chat_replay': YoutubeLiveChatReplayFD,
 }
 
@@ -52,6 +54,7 @@ def shorten_protocol_name(proto, simplify=False):
         'rtmp_ffmpeg': 'rtmp_f',
         'http_dash_segments': 'dash',
         'niconico_dmc': 'dmc',
+        'websocket_frag': 'WSfrag',
     }
     if simplify:
         short_protocol_names.update({

yt_dlp/downloader/external.py

@@ -347,6 +347,10 @@ class FFmpegFD(ExternalFD):
         # TODO: Fix path for ffmpeg
         return FFmpegPostProcessor().available
 
+    def on_process_started(self, proc, stdin):
+        """ Override this in subclasses  """
+        pass
+
     def _call_downloader(self, tmpfilename, info_dict):
         urls = [f['url'] for f in info_dict.get('requested_formats', [])] or [info_dict['url']]
         ffpp = FFmpegPostProcessor(downloader=self)
@@ -474,6 +478,8 @@ class FFmpegFD(ExternalFD):
         self._debug_cmd(args)
 
         proc = subprocess.Popen(args, stdin=subprocess.PIPE, env=env)
+        if url in ('-', 'pipe:'):
+            self.on_process_started(proc, proc.stdin)
         try:
             retval = proc.wait()
         except BaseException as e:
@@ -482,7 +488,7 @@ class FFmpegFD(ExternalFD):
             # produces a file that is playable (this is mostly useful for live
             # streams). Note that Windows is not affected and produces playable
             # files (see https://github.com/ytdl-org/youtube-dl/issues/8300).
-            if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32':
+            if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32' and url not in ('-', 'pipe:'):
                 process_communicate_or_kill(proc, b'q')
             else:
                 proc.kill()

yt_dlp/downloader/websocket.py

@@ -0,0 +1,59 @@
+import os
+import signal
+import asyncio
+import threading
+
+try:
+    import websockets
+    has_websockets = True
+except ImportError:
+    has_websockets = False
+
+from .common import FileDownloader
+from .external import FFmpegFD
+
+
+class FFmpegSinkFD(FileDownloader):
+    """ A sink to ffmpeg for downloading fragments in any form """
+
+    def real_download(self, filename, info_dict):
+        info_copy = info_dict.copy()
+        info_copy['url'] = '-'
+
+        async def call_conn(proc, stdin):
+            try:
+                await self.real_connection(stdin, info_dict)
+            except (BrokenPipeError, OSError):
+                pass
+            finally:
+                try:
+                    stdin.flush()
+                    stdin.close()
+                except OSError:
+                    pass
+                os.kill(os.getpid(), signal.SIGINT)
+
+        class FFmpegStdinFD(FFmpegFD):
+            @classmethod
+            def get_basename(cls):
+                return FFmpegFD.get_basename()
+
+            def on_process_started(self, proc, stdin):
+                thread = threading.Thread(target=asyncio.run, daemon=True, args=(call_conn(proc, stdin), ))
+                thread.start()
+
+        return FFmpegStdinFD(self.ydl, self.params or {}).download(filename, info_copy)
+
+    async def real_connection(self, sink, info_dict):
+        """ Override this in subclasses """
+        raise NotImplementedError('This method must be implemented by subclasses')
+
+
+class WebSocketFragmentFD(FFmpegSinkFD):
+    async def real_connection(self, sink, info_dict):
+        async with websockets.connect(info_dict['url'], extra_headers=info_dict.get('http_headers', {})) as ws:
+            while True:
+                recv = await ws.recv()
+                if isinstance(recv, str):
+                    recv = recv.encode('utf8')
+                sink.write(recv)

yt_dlp/extractor/common.py

@@ -1487,7 +1487,7 @@ class InfoExtractor(object):
             'acodec': {'type': 'ordered', 'regex': True,
                        'order': ['opus', 'vorbis', 'aac', 'mp?4a?', 'mp3', 'e?a?c-?3', 'dts', '', None, 'none']},
             'proto': {'type': 'ordered', 'regex': True, 'field': 'protocol',
-                      'order': ['(ht|f)tps', '(ht|f)tp$', 'm3u8.+', 'm3u8', '.*dash', '', 'mms|rtsp', 'none', 'f4']},
+                      'order': ['(ht|f)tps', '(ht|f)tp$', 'm3u8.+', '.*dash', 'ws|websocket', '', 'mms|rtsp', 'none', 'f4']},
             'vext': {'type': 'ordered', 'field': 'video_ext',
                      'order': ('mp4', 'webm', 'flv', '', 'none'),
                      'order_free': ('webm', 'mp4', 'flv', '', 'none')},

yt_dlp/options.py

@@ -1165,7 +1165,7 @@ def parseOpts(overrideArguments=None):
             'to give the argument to the specified postprocessor/executable. Supported PP are: '
             'Merger, ExtractAudio, SplitChapters, Metadata, EmbedSubtitle, EmbedThumbnail, '
             'SubtitlesConvertor, ThumbnailsConvertor, VideoRemuxer, VideoConvertor, '
-            'SponSkrub, FixupStretched, FixupM4a and FixupM3u8. '
+            'SponSkrub, FixupStretched, FixupM4a, FixupM3u8, FixupTimestamp and FixupDuration. '
             'The supported executables are: AtomicParsley, FFmpeg, FFprobe, and SponSkrub. '
             'You can also specify "PP+EXE:ARGS" to give the arguments to the specified executable '
             'only when being used by the specified postprocessor. Additionally, for ffmpeg/ffprobe, '

yt_dlp/postprocessor/__init__.py

@@ -5,7 +5,9 @@ from .ffmpeg import (
     FFmpegPostProcessor,
     FFmpegEmbedSubtitlePP,
     FFmpegExtractAudioPP,
+    FFmpegFixupDurationPP,
     FFmpegFixupStretchedPP,
+    FFmpegFixupTimestampPP,
     FFmpegFixupM3u8PP,
     FFmpegFixupM4aPP,
     FFmpegMergerPP,
@@ -35,9 +37,11 @@ __all__ = [
     'FFmpegEmbedSubtitlePP',
     'FFmpegExtractAudioPP',
     'FFmpegSplitChaptersPP',
+    'FFmpegFixupDurationPP',
     'FFmpegFixupM3u8PP',
     'FFmpegFixupM4aPP',
     'FFmpegFixupStretchedPP',
+    'FFmpegFixupTimestampPP',
     'FFmpegMergerPP',
     'FFmpegMetadataPP',
     'FFmpegSubtitlesConvertorPP',

yt_dlp/postprocessor/ffmpeg.py

@@ -700,6 +700,35 @@ class FFmpegFixupM3u8PP(FFmpegFixupPostProcessor):
         return [], info
 
 
+class FFmpegFixupTimestampPP(FFmpegFixupPostProcessor):
+
+    def __init__(self, downloader=None, trim=0.001):
+        # "trim" should be used when the video contains unintended packets
+        super(FFmpegFixupTimestampPP, self).__init__(downloader)
+        assert isinstance(trim, (int, float))
+        self.trim = str(trim)
+
+    @PostProcessor._restrict_to(images=False)
+    def run(self, info):
+        required_version = '4.4'
+        if is_outdated_version(self._versions[self.basename], required_version):
+            self.report_warning(
+                'A re-encode is needed to fix timestamps in older versions of ffmpeg. '
+                f'Please install ffmpeg {required_version} or later to fixup without re-encoding')
+            opts = ['-vf', 'setpts=PTS-STARTPTS']
+        else:
+            opts = ['-c', 'copy', '-bsf', 'setts=ts=TS-STARTPTS']
+        self._fixup('Fixing frame timestamp', info['filepath'], opts + ['-map', '0', '-dn', '-ss', self.trim])
+        return [], info
+
+
+class FFmpegFixupDurationPP(FFmpegFixupPostProcessor):
+    @PostProcessor._restrict_to(images=False)
+    def run(self, info):
+        self._fixup('Fixing video duration', info['filepath'], ['-c', 'copy', '-map', '0', '-dn'])
+        return [], info
+
+
 class FFmpegSubtitlesConvertorPP(FFmpegPostProcessor):
     SUPPORTED_EXTS = ('srt', 'vtt', 'ass', 'lrc')