From 9bba470bd1c4684d40440d9b5ca7d97944e63d70 Mon Sep 17 00:00:00 2001
From: Hellow <74311245+HeIIow2@users.noreply.github.com>
Date: Thu, 15 Jun 2023 18:22:00 +0200
Subject: [PATCH] connection

---
 src/music_kraken/connection/connection.py | 52 ++++++++++++++++++++++-
 src/music_kraken/objects/target.py        |  5 ++-
 src/music_kraken/pages/youtube.py         | 13 +++---
 3 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/src/music_kraken/connection/connection.py b/src/music_kraken/connection/connection.py
index 46c22de..d243bb4 100644
--- a/src/music_kraken/connection/connection.py
+++ b/src/music_kraken/connection/connection.py
@@ -4,9 +4,11 @@ from urllib.parse import urlparse, urlunsplit, ParseResult
 import logging
 
 import requests
+from tqdm import tqdm
 
 from .rotating import RotatingProxy
-from ..utils.shared import PROXIES_LIST
+from ..utils.shared import PROXIES_LIST, CHUNK_SIZE
+from ..objects import Target
 
 
 class Connection:
@@ -196,3 +198,51 @@ class Connection:
             self.LOGGER.warning(f"Max attempts ({self.TRIES}) exceeded for: GET:{url}")
             self.LOGGER.warning(f"payload: {json}")
         return r
+
+    def stream_into(
+            self,
+            url: str,
+            target: Target,
+            description: str = "download",
+            refer_from_origin: bool = True,
+            accepted_response_codes: set = None,
+            timeout: float = None,
+            headers: dict = None,
+            raw_url: bool = False,
+            **kwargs
+    ):
+        r = self._request(
+            request=self.session.get,
+            try_count=0,
+            accepted_response_code=accepted_response_codes or self.ACCEPTED_RESPONSE_CODES,
+            url=url,
+            timeout=timeout,
+            headers=headers,
+            raw_url=raw_url,
+            refer_from_origin=refer_from_origin,
+            stream=True,
+            **kwargs
+        )
+
+        if r is None:
+            return False
+
+        target.create_path()
+        total_size = int(r.headers.get('content-length'))
+
+        with target.open("wb") as f:
+            try:
+                """
+                https://en.wikipedia.org/wiki/Kilobyte
+                > The internationally recommended unit symbol for the kilobyte is kB.
+                """
+                with tqdm(total=total_size, unit='B', unit_scale=True, unit_divisor=1024, desc=description) as t:
+
+                    for chunk in r.iter_content(chunk_size=CHUNK_SIZE):
+                        size = f.write(chunk)
+                        t.update(size)
+                return True
+
+            except requests.exceptions.ConnectionError:
+                self.LOGGER.error("Stream timed out.")
+                return False
diff --git a/src/music_kraken/objects/target.py b/src/music_kraken/objects/target.py
index 8c70755..b979951 100644
--- a/src/music_kraken/objects/target.py
+++ b/src/music_kraken/objects/target.py
@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import List, Tuple
+from typing import List, Tuple, TextIO
 
 import requests
 from tqdm import tqdm
@@ -98,6 +98,9 @@ class Target(DatabaseObject):
             except requests.exceptions.Timeout:
                 shared.DOWNLOAD_LOGGER.error("Stream timed out.")
                 return False
+
+    def open(self, file_mode: str, **kwargs) -> TextIO:
+        return self.file_path.open(file_mode, **kwargs)
             
     def delete(self):
         self.file_path.unlink(missing_ok=True)
diff --git a/src/music_kraken/pages/youtube.py b/src/music_kraken/pages/youtube.py
index 4e5b9fe..797480e 100644
--- a/src/music_kraken/pages/youtube.py
+++ b/src/music_kraken/pages/youtube.py
@@ -5,8 +5,6 @@ from enum import Enum
 import sponsorblock
 from sponsorblock.errors import HTTPException, NotFoundException
 
-from music_kraken.objects import Song, Source
-
 from ..objects import Source, DatabaseObject, Song, Target
 from .abstract import Page
 from ..objects import (
@@ -373,11 +371,9 @@ class YouTube(Page):
 
         endpoint = audio_format["url"]
 
-        r = self.download_connection.get(endpoint, stream=True, raw_url=True)
-        if r is None:
-            return DownloadResult(error_message=f"Couldn't connect to {endpoint}")
+        self.download_connection.stream_into(endpoint, target, description=desc, raw_url=True)
 
-        if target.stream_into(r, desc=desc):
+        if self.download_connection.get(endpoint, stream=True, raw_url=True):
             return DownloadResult(total=1)
         return DownloadResult(error_message=f"Streaming to the file went wrong: {endpoint}, {str(target.file_path)}")
 
@@ -397,5 +393,8 @@ class YouTube(Page):
             self.LOGGER.debug(f"No sponsor found for the video {parsed.id}.")
         except HTTPException as e:
             self.LOGGER.warning(f"{e}")
-            
+
+        if len(segments) > 0:
+            print(f"Removing {len(segments)} interruptions in the audio...")
+
         return [(segment.start, segment.end) for segment in segments]