diff --git a/pyproject.toml b/pyproject.toml index cd553f4..e3db227 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,6 @@ classifiers = [ dependencies = [ "click>=8.0.4", "requests>=2.31.0", - "wget>=3.2", ] [project.urls] diff --git a/src/apyrat.py b/src/apyrat.py index 563f886..da4f679 100644 --- a/src/apyrat.py +++ b/src/apyrat.py @@ -1,12 +1,12 @@ import hashlib import os +import tempfile from concurrent.futures import ThreadPoolExecutor, as_completed from enum import Enum from urllib.parse import urlparse import click import requests -import wget from src.utils import check_domain_validity, prepare_headers @@ -142,20 +142,47 @@ def download(self, quality: str): video = self.find_closest_quality(video_qualities, quality) self._download_video(video) + @staticmethod + def _sanitize_filename(name: str) -> str: + # Remove characters that are invalid on common filesystems + invalid = r'<>:"/\|?*' + for char in invalid: + name = name.replace(char, "") + # Strip trailing dots/spaces and limit length to be safe + name = name.strip().strip(".") + if len(name) > 200: + name = name[:200] + return name or "video" + def _download_video(self, video): url = video.get("url") title = video.get("title") click.echo(f"{title}\n{url}") - file_name = self.file_name or video.get("title") - output_file = f"{file_name}.{self._get_file_format(url)}" + file_name = self._sanitize_filename(self.file_name or video.get("title")) + file_format = self._get_file_format(url) + output_file = f"{file_name}.{file_format}" if file_format else file_name # Check if file already exists if os.path.isfile(output_file): # Create a hash of the filename hash_object = hashlib.md5(output_file.encode()) hex_dig = hash_object.hexdigest() # Append the hash to the filename - output_file = f"{hex_dig}{output_file}" - wget.download(url, out=output_file) + output_file = f"{hex_dig}_{output_file}" + # Download to a safe temporary file first, then rename. This avoids + # issues with filesystems/locales that choke on long Unicode paths + # (e.g. Persian filenames) when opening the final file directly. + fd, tmp_file = tempfile.mkstemp(dir=".") + os.close(fd) + try: + with requests.get(url, headers=prepare_headers(), stream=True) as r: + r.raise_for_status() + with open(tmp_file, "wb") as f: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + os.replace(tmp_file, output_file) + finally: + if os.path.exists(tmp_file): + os.remove(tmp_file) @staticmethod def _get_file_format(url): diff --git a/tests/test_apyrat.py b/tests/test_apyrat.py index 0612f1a..b1e5e40 100644 --- a/tests/test_apyrat.py +++ b/tests/test_apyrat.py @@ -1,4 +1,6 @@ -from unittest.mock import patch +import hashlib +import os +from unittest.mock import MagicMock, patch from src.apyrat import Downloader, VideoQuality @@ -29,8 +31,13 @@ def test_default_quality(): @patch("os.path.isfile", return_value=True) -@patch("wget.download") -def test_download_when_file_exists(mock_wget, mock_isfile): +@patch("requests.get") +def test_download_when_file_exists(mock_get, mock_isfile): + mock_response = MagicMock() + mock_response.iter_content.return_value = [b"chunk"] + mock_response.raise_for_status = MagicMock() + mock_get.return_value.__enter__ = MagicMock(return_value=mock_response) + mock_get.return_value.__exit__ = MagicMock(return_value=False) downloader = Downloader("https://www.aparat.com/v/qur3I", "outputfile") downloader.videos = [ [ @@ -42,6 +49,12 @@ def test_download_when_file_exists(mock_wget, mock_isfile): ] ] downloader.qualities = ["720p"] - downloader.download("720p") - mock_wget.assert_called_once() + # The URL has no extension, so the output file is just "outputfile" + hex_dig = hashlib.md5("outputfile".encode()).hexdigest() + output_file = f"{hex_dig}_outputfile" + try: + downloader.download("720p") + finally: + if os.path.exists(output_file): + os.remove(output_file) mock_isfile.assert_called_once()