HF_Agents_Final_Project / src /download_utils.py
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.")