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 logging
import requests import requests
from tqdm import tqdm
from .rotating import RotatingProxy from .rotating import RotatingProxy
from ..utils.shared import PROXIES_LIST from ..utils.shared import PROXIES_LIST, CHUNK_SIZE
from ..objects import Target
class Connection: class Connection:
@ -196,3 +198,51 @@ class Connection:
self.LOGGER.warning(f"Max attempts ({self.TRIES}) exceeded for: GET:{url}") self.LOGGER.warning(f"Max attempts ({self.TRIES}) exceeded for: GET:{url}")
self.LOGGER.warning(f"payload: {json}") self.LOGGER.warning(f"payload: {json}")
return r 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 pathlib import Path
from typing import List, Tuple from typing import List, Tuple, TextIO
import requests import requests
from tqdm import tqdm from tqdm import tqdm
@ -99,5 +99,8 @@ class Target(DatabaseObject):
shared.DOWNLOAD_LOGGER.error("Stream timed out.") shared.DOWNLOAD_LOGGER.error("Stream timed out.")
return False return False
def open(self, file_mode: str, **kwargs) -> TextIO:
return self.file_path.open(file_mode, **kwargs)
def delete(self): def delete(self):
self.file_path.unlink(missing_ok=True) self.file_path.unlink(missing_ok=True)

View File

@ -5,8 +5,6 @@ from enum import Enum
import sponsorblock import sponsorblock
from sponsorblock.errors import HTTPException, NotFoundException from sponsorblock.errors import HTTPException, NotFoundException
from music_kraken.objects import Song, Source
from ..objects import Source, DatabaseObject, Song, Target from ..objects import Source, DatabaseObject, Song, Target
from .abstract import Page from .abstract import Page
from ..objects import ( from ..objects import (
@ -373,11 +371,9 @@ class YouTube(Page):
endpoint = audio_format["url"] endpoint = audio_format["url"]
r = self.download_connection.get(endpoint, stream=True, raw_url=True) self.download_connection.stream_into(endpoint, target, description=desc, raw_url=True)
if r is None:
return DownloadResult(error_message=f"Couldn't connect to {endpoint}")
if target.stream_into(r, desc=desc): if self.download_connection.get(endpoint, stream=True, raw_url=True):
return DownloadResult(total=1) return DownloadResult(total=1)
return DownloadResult(error_message=f"Streaming to the file went wrong: {endpoint}, {str(target.file_path)}") 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: except HTTPException as e:
self.LOGGER.warning(f"{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] return [(segment.start, segment.end) for segment in segments]