[downloader/fragment] Allow persisting extra state when a download is interrupted

1 files changed, 6 insertions, 1 deletions

diff --git a/yt_dlp/downloader/fragment.py b/yt_dlp/downloader/fragment.py
index a0c1d13ac..fadd0dfc5 100644
--- a/yt_dlp/downloader/fragment.py
+++ b/yt_dlp/downloader/fragment.py
@@ -77,7 +77,10 @@ class FragmentFD(FileDownloader):
         assert 'ytdl_corrupt' not in ctx
         stream, _ = sanitize_open(self.ytdl_filename(ctx['filename']), 'r')
         try:
-            ctx['fragment_index'] = json.loads(stream.read())['downloader']['current_fragment']['index']
+            ytdl_data = json.loads(stream.read())
+            ctx['fragment_index'] = ytdl_data['downloader']['current_fragment']['index']
+            if 'extra_state' in ytdl_data['downloader']:
+                ctx['extra_state'] = ytdl_data['downloader']['extra_state']
         except Exception:
             ctx['ytdl_corrupt'] = True
         finally:
@@ -90,6 +93,8 @@ class FragmentFD(FileDownloader):
                 'index': ctx['fragment_index'],
             },
         }
+        if 'extra_state' in ctx:
+            downloader['extra_state'] = ctx['extra_state']
         if ctx.get('fragment_count') is not None:
             downloader['fragment_count'] = ctx['fragment_count']
         frag_index_stream.write(json.dumps({'downloader': downloader}))