From 4d6c72ba307b366b3df72cc5243bf8b6a593de11 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Sat, 9 Mar 2024 01:37:55 +0000
Subject: [PATCH] fix(DASH/HLS): Don't merge folders, skip final merge if only
 1 segment

---
 devine/core/manifests/dash.py | 43 +++++++++++++++++++----------------
 devine/core/manifests/hls.py  | 36 +++++++++++++++++------------
 2 files changed, 45 insertions(+), 34 deletions(-)

diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py
index 00db1e0..72f5d26 100644
--- a/devine/core/manifests/dash.py
+++ b/devine/core/manifests/dash.py
@@ -486,28 +486,33 @@ class DASH:
         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
index 60046e8..aab72c8 100644
--- a/devine/core/manifests/hls.py
+++ b/devine/core/manifests/hls.py
@@ -508,22 +508,28 @@ class HLS:
             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")