Yago Bolivar
feat: implement file download utility with error handling and filename inference
c511b4a
import requests | |
import os | |
import shutil | |
def download_file(url: str, save_dir: str, filename: str = None) -> str | None: | |
""" | |
Downloads a file from a URL and saves it to the specified directory. | |
Args: | |
url: The URL of the file to download. | |
save_dir: The directory where the file should be saved. | |
filename: Optional. The name to save the file as. | |
If None, tries to infer from URL or Content-Disposition. | |
Returns: | |
The full path to the downloaded file if successful, None otherwise. | |
""" | |
if not url: | |
print("Error: Download URL cannot be empty.") | |
return None | |
os.makedirs(save_dir, exist_ok=True) | |
try: | |
with requests.get(url, stream=True, timeout=30) as r: | |
r.raise_for_status() # Raise an exception for bad status codes | |
if not filename: | |
# Try to get filename from Content-Disposition header | |
content_disposition = r.headers.get('content-disposition') | |
if content_disposition: | |
import re | |
fname_match = re.findall('filename="?([^"]+)"?', content_disposition) | |
if fname_match: | |
filename = fname_match[0] | |
if not filename: # Fallback to last part of URL | |
filename = url.split('/')[-1] | |
if not filename: # If URL ends with '/', generate a name | |
filename = "downloaded_file" | |
# Sanitize filename (basic example, might need more robust sanitization) | |
filename = "".join(c for c in filename if c.isalnum() or c in ['.', '_', '-']).strip() | |
if not filename: # If sanitization results in empty filename | |
filename = "downloaded_file_unnamed" | |
local_filepath = os.path.join(save_dir, filename) | |
with open(local_filepath, 'wb') as f: | |
shutil.copyfileobj(r.raw, f) | |
print(f"Successfully downloaded '{filename}' to '{local_filepath}'") | |
return local_filepath | |
except requests.exceptions.RequestException as e: | |
print(f"Error downloading file from {url}: {e}") | |
return None | |
except IOError as e: | |
print(f"Error saving file to {local_filepath}: {e}") | |
return None | |
except Exception as e: | |
print(f"An unexpected error occurred during download: {e}") | |
return None | |
if __name__ == '__main__': | |
# Example Usage: | |
test_url = "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf" # A sample PDF | |
download_dir = "test_downloads" | |
# Test 1: Basic download | |
print("\n--- Test 1: Basic Download ---") | |
downloaded_path = download_file(test_url, download_dir) | |
if downloaded_path and os.path.exists(downloaded_path): | |
print(f"Test 1 Success: File at {downloaded_path}") | |
else: | |
print("Test 1 Failed.") | |
# Test 2: Download with specified filename | |
print("\n--- Test 2: Download with specified filename ---") | |
custom_filename = "my_custom_dummy.pdf" | |
downloaded_path_custom = download_file(test_url, download_dir, filename=custom_filename) | |
if downloaded_path_custom and os.path.exists(downloaded_path_custom) and os.path.basename(downloaded_path_custom) == custom_filename: | |
print(f"Test 2 Success: File at {downloaded_path_custom}") | |
else: | |
print("Test 2 Failed.") | |
# Test 3: Invalid URL | |
print("\n--- Test 3: Invalid URL ---") | |
invalid_url = "http://invalid.url/nonexistentfile.txt" | |
downloaded_path_invalid = download_file(invalid_url, download_dir) | |
if downloaded_path_invalid is None: | |
print("Test 3 Success: Handled invalid URL correctly.") | |
else: | |
print("Test 3 Failed.") | |
# Cleanup (optional) | |
# if os.path.exists(download_dir): | |
# import shutil | |
# shutil.rmtree(download_dir) | |
# print(f"\nCleaned up '{download_dir}' directory.") |