connection

This commit is contained in:
Hellow 2023-06-15 18:22:00 +02:00
parent 6197345023
commit 9bba470bd1
3 changed files with 61 additions and 9 deletions

View File

@ -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

View File

@ -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
@ -99,5 +99,8 @@ class Target(DatabaseObject):
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)

View File

@ -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)}")
@ -398,4 +394,7 @@ class YouTube(Page):
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]