fix(DASH/HLS): Don't merge folders, skip final merge if only 1 segment

devine-dl · Mar 9, 2024 · 4d6c72b · 4d6c72b
1 parent 77e663e
commit 4d6c72b
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 34 deletions.
diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py
@@ -486,28 +486,33 @@ def download_track(
         for control_file in save_dir.glob("*.aria2__temp"):
             control_file.unlink()
 
-        segments_to_merge = sorted(save_dir.iterdir())
-        progress(downloaded="Merging", completed=0, total=len(segments_to_merge))
-
+        segments_to_merge = [
+            x
+            for x in sorted(save_dir.iterdir())
+            if x.is_file()
+        ]
         with open(save_path, "wb") as f:
             if init_data:
                 f.write(init_data)
-            for segment_file in segments_to_merge:
-                segment_data = segment_file.read_bytes()
-                # TODO: fix encoding after decryption?
-                if (
-                    not drm and isinstance(track, Subtitle) and
-                    track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML)
-                ):
-                    segment_data = try_ensure_utf8(segment_data)
-                    segment_data = segment_data.decode("utf8"). \
-                        replace("&lrm;", html.unescape("&lrm;")). \
-                        replace("&rlm;", html.unescape("&rlm;")). \
-                        encode("utf8")
-                f.write(segment_data)
-                f.flush()
-                segment_file.unlink()
-                progress(advance=1)
+            if len(segments_to_merge) > 1:
+                progress(downloaded="Merging", completed=0, total=len(segments_to_merge))
+            else:
+                for segment_file in segments_to_merge:
+                    segment_data = segment_file.read_bytes()
+                    # TODO: fix encoding after decryption?
+                    if (
+                        not drm and isinstance(track, Subtitle) and
+                        track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML)
+                    ):
+                        segment_data = try_ensure_utf8(segment_data)
+                        segment_data = segment_data.decode("utf8"). \
+                            replace("&lrm;", html.unescape("&lrm;")). \
+                            replace("&rlm;", html.unescape("&rlm;")). \
+                            encode("utf8")
+                    f.write(segment_data)
+                    f.flush()
+                    segment_file.unlink()
+                    progress(advance=1)
 
         track.path = save_path
         if callable(track.OnDownloaded):

diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py
@@ -508,22 +508,28 @@ def merge_discontinuity(include_this_segment: bool, include_map_data: bool = Tru
             return
 
         # finally merge all the discontinuity save files together to the final path
-        progress(downloaded="Merging")
-        if isinstance(track, (Video, Audio)):
-            HLS.merge_segments(
-                segments=sorted(list(save_dir.iterdir())),
-                save_path=save_path
-            )
-            shutil.rmtree(save_dir)
+        segments_to_merge = [
+            x
+            for x in sorted(save_dir.iterdir())
+            if x.is_file()
+        ]
+        if len(segments_to_merge) == 1:
+            shutil.move(segments_to_merge[0], save_path)
         else:
-            with open(save_path, "wb") as f:
-                for discontinuity_file in sorted(save_dir.iterdir()):
-                    if discontinuity_file.is_dir():
-                        continue
-                    discontinuity_data = discontinuity_file.read_bytes()
-                    f.write(discontinuity_data)
-                    f.flush()
-            shutil.rmtree(save_dir)
+            progress(downloaded="Merging")
+            if isinstance(track, (Video, Audio)):
+                HLS.merge_segments(
+                    segments=segments_to_merge,
+                    save_path=save_path
+                )
+                shutil.rmtree(save_dir)
+            else:
+                with open(save_path, "wb") as f:
+                    for discontinuity_file in segments_to_merge:
+                        discontinuity_data = discontinuity_file.read_bytes()
+                        f.write(discontinuity_data)
+                        f.flush()
+                shutil.rmtree(save_dir)
 
         progress(downloaded="Downloaded")