From 840db6e6896cbc5c40745acf8a24ecafad3f9f70 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Wed, 1 Mar 2023 08:54:35 +0000
Subject: [PATCH] Move segment merging from dl to DASH/HLS classes

---
 devine/commands/dl.py         |  9 ++-------
 devine/core/manifests/dash.py | 25 ++++++++++++++++++++-----
 devine/core/manifests/hls.py  | 25 ++++++++++++++++++++-----
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/devine/commands/dl.py b/devine/commands/dl.py
index a760ad9..5af9eea 100644
--- a/devine/commands/dl.py
+++ b/devine/commands/dl.py
@@ -735,6 +735,7 @@ class dl:
         if track.descriptor == track.Descriptor.M3U:
             HLS.download_track(
                 track=track,
+                save_path=save_path,
                 save_dir=save_dir,
                 stop_event=self.DL_POOL_STOP,
                 progress=progress,
@@ -745,6 +746,7 @@ class dl:
         elif track.descriptor == track.Descriptor.MPD:
             DASH.download_track(
                 track=track,
+                save_path=save_path,
                 save_dir=save_dir,
                 stop_event=self.DL_POOL_STOP,
                 progress=progress,
@@ -791,13 +793,6 @@ class dl:
                 track.drm = None
                 if callable(track.OnDecrypted):
                     track.OnDecrypted(track)
-        else:
-            with open(save_path, "wb") as f:
-                for file in sorted(save_dir.iterdir()):
-                    f.write(file.read_bytes())
-                    file.unlink()
-            save_dir.rmdir()
-            track.path = save_path
 
         if track.path.stat().st_size <= 3:  # Empty UTF-8 BOM == 3 bytes
             raise IOError(
diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py
index 5295d97..7a104a8 100644
--- a/devine/core/manifests/dash.py
+++ b/devine/core/manifests/dash.py
@@ -4,6 +4,7 @@ import base64
 import logging
 import math
 import re
+import shutil
 import sys
 import time
 import traceback
@@ -272,6 +273,7 @@ class DASH:
     @staticmethod
     def download_track(
         track: AnyTrack,
+        save_path: Path,
         save_dir: Path,
         stop_event: Event,
         progress: partial,
@@ -504,13 +506,13 @@ class DASH:
 
             progress(total=len(segments))
 
+            finished_threads = 0
+            has_stopped = False
+            has_failed = False
             download_sizes = []
             last_speed_refresh = time.time()
 
             with ThreadPoolExecutor(max_workers=16) as pool:
-                finished_threads = 0
-                has_stopped = False
-                has_failed = False
                 for download in futures.as_completed((
                     pool.submit(
                         download_segment,
@@ -556,10 +558,23 @@ class DASH:
                         progress(downloaded=f"DASH {filesize.decimal(download_speed)}/s")
                         last_speed_refresh = now
                         download_sizes.clear()
-                if has_failed:
-                    progress(downloaded="[red]FAILED")
+
+            try:
                 if has_stopped:
                     progress(downloaded="[yellow]STOPPED")
+                    return
+                if has_failed:
+                    progress(downloaded="[red]FAILED")
+                    return
+
+                with open(save_path, "wb") as f:
+                    for segment_file in sorted(save_dir.iterdir()):
+                        f.write(segment_file.read_bytes())
+                        segment_file.unlink()
+
+                track.path = save_path
+            finally:
+                shutil.rmtree(save_dir)
 
     @staticmethod
     def get_language(*options: Any) -> Optional[Language]:
diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py
index 6d8eeaf..ed99f7d 100644
--- a/devine/core/manifests/hls.py
+++ b/devine/core/manifests/hls.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import logging
 import re
+import shutil
 import sys
 import time
 import traceback
@@ -181,6 +182,7 @@ class HLS:
     @staticmethod
     def download_track(
         track: AnyTrack,
+        save_path: Path,
         save_dir: Path,
         stop_event: Event,
         progress: partial,
@@ -342,13 +344,13 @@ class HLS:
 
         progress(total=len(master.segments))
 
+        finished_threads = 0
+        has_stopped = False
+        has_failed = False
         download_sizes = []
         last_speed_refresh = time.time()
 
         with ThreadPoolExecutor(max_workers=16) as pool:
-            finished_threads = 0
-            has_stopped = False
-            has_failed = False
             for download in futures.as_completed((
                 pool.submit(
                     download_segment,
@@ -396,10 +398,23 @@ class HLS:
                     progress(downloaded=f"HLS {filesize.decimal(download_speed)}/s")
                     last_speed_refresh = now
                     download_sizes.clear()
-            if has_failed:
-                progress(downloaded="[red]FAILED")
+
+        try:
             if has_stopped:
                 progress(downloaded="[yellow]STOPPED")
+                return
+            if has_failed:
+                progress(downloaded="[red]FAILED")
+                return
+
+            with open(save_path, "wb") as f:
+                for segment_file in sorted(save_dir.iterdir()):
+                    f.write(segment_file.read_bytes())
+                    segment_file.unlink()
+
+            track.path = save_path
+        finally:
+            shutil.rmtree(save_dir)
 
     @staticmethod
     def get_drm(