From ac7cdba28e1798045a7fab3b59365464fe30263c Mon Sep 17 00:00:00 2001
From: Ivan Habunek <ivan@habunek.com>
Date: Fri, 30 Aug 2024 11:42:54 +0200
Subject: [PATCH] Download init segments

These seem to occur in hd quality playlists like 1440p.
---
 twitchdl/commands/download.py |  6 ++++++
 twitchdl/playlists.py         | 11 ++++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/twitchdl/commands/download.py b/twitchdl/commands/download.py
index 3bdcd45..689f107 100644
--- a/twitchdl/commands/download.py
+++ b/twitchdl/commands/download.py
@@ -19,6 +19,7 @@ from twitchdl.naming import clip_filename, video_filename
 from twitchdl.output import blue, bold, green, print_log, yellow
 from twitchdl.playlists import (
     enumerate_vods,
+    get_init_sections,
     load_m3u8,
     make_join_playlist,
     parse_playlists,
@@ -231,6 +232,11 @@ def _download_video(video_id: str, args: DownloadOptions) -> None:
 
     click.echo(f"\nDownloading {len(vods)} VODs using {args.max_workers} workers to {target_dir}")
 
+    init_sections = get_init_sections(vods_m3u8)
+    for uri in init_sections:
+        print_log(f"Downloading init section {uri}...")
+        download_file(f"{base_uri}{uri}", target_dir / uri)
+
     sources = [base_uri + vod.path for vod in vods]
     targets = [target_dir / f"{vod.index:05d}.ts" for vod in vods]
 
diff --git a/twitchdl/playlists.py b/twitchdl/playlists.py
index 9564827..4570e5a 100644
--- a/twitchdl/playlists.py
+++ b/twitchdl/playlists.py
@@ -4,7 +4,7 @@ Parse and manipulate m3u8 playlists.
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Generator, List, Optional, OrderedDict
+from typing import Generator, List, Optional, OrderedDict, Set
 
 import click
 import m3u8
@@ -169,3 +169,12 @@ def _playlist_key(playlist: Playlist) -> int:
         pass
 
     return MAX
+
+
+def get_init_sections(playlist: m3u8.M3U8) -> Set[str]:
+    # TODO: we're ignoring initi_section.base_uri and bytes
+    return set(
+        segment.init_section.uri
+        for segment in playlist.segments
+        if segment.init_section is not None
+    )