Update main.py

This commit is contained in:
nloginov 2024-07-06 12:35:02 +02:00 committed by GitHub
parent 528ea0848f
commit 19da1d5c0e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

45
main.py
View file

@ -64,6 +64,34 @@ def extract_and_fix_urls(text):
return fixed_urls
def get_imgur_url(imgur_url):
try:
# Handle album URLs
if '/a/' in imgur_url:
response = requests.get(imgur_url, timeout=10)
response.raise_for_status()
# Extract all image URLs from the album page
image_urls = re.findall(r'https://i\.imgur\.com/\w+\.(?:jpg|png|gif|mp4)', response.text)
return image_urls if image_urls else None
# Handle single image/video URLs
response = requests.get(imgur_url, timeout=10)
response.raise_for_status()
content = response.text
# Try to find direct image/video URL in the page source
match = re.search(r'https://i\.imgur\.com/\w+\.(?:jpg|png|gif|mp4)', content)
if match:
return match.group(0)
# If direct URL not found, construct it from the imgur ID
imgur_id = imgur_url.split('/')[-1]
return f'https://i.imgur.com/{imgur_id}.jpg'
except requests.exceptions.RequestException:
pass # Silently handle the error
return None
def get_tenor_gif_url(tenor_url):
try:
response = requests.get(tenor_url, timeout=10)
@ -115,8 +143,19 @@ def safe_filename(filename, max_length=200):
def download_media(url):
global successful_downloads, failed_downloads
try:
if url.lower().endswith(SUPPORTED_EXTENSIONS):
direct_url = url
if 'imgur.com' in url:
imgur_urls = get_imgur_url(url)
if imgur_urls:
if isinstance(imgur_urls, list): # It's an album
for imgur_url in imgur_urls:
download_media(imgur_url) # Recursive call for each image in the album
return
else: # Single image/video
direct_url = imgur_urls
else:
failed_downloads += 1
error_summary["Imgur URL skipped"].append(url)
return
elif 'tenor.com' in url:
gif_url = get_tenor_gif_url(url)
if gif_url:
@ -125,6 +164,8 @@ def download_media(url):
failed_downloads += 1
error_summary["Tenor URL skipped"].append(url)
return
elif url.lower().endswith(SUPPORTED_EXTENSIONS):
direct_url = url
else:
direct_url = url