Improve logo processing: Add image verification and better error handling

This commit is contained in:
ChuckBuilds
2025-04-11 13:05:16 -05:00
parent b899f56cab
commit c9c385b2aa

View File

@@ -377,21 +377,60 @@ class StockManager:
# Try Yahoo first # Try Yahoo first
response = requests.get(yahoo_url, headers=self.headers, timeout=5) response = requests.get(yahoo_url, headers=self.headers, timeout=5)
if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''): if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
# Create image from response content try:
from io import BytesIO # Create image from response content
logo = Image.open(BytesIO(response.content)) from io import BytesIO
logo = Image.open(BytesIO(response.content))
# Verify it's a valid image
logo.verify()
# Reopen after verify
logo = Image.open(BytesIO(response.content))
except Exception as e:
logger.warning(f"Invalid image data from Yahoo for {symbol}: {str(e)}")
# Try alternative source
response = requests.get(alt_url, headers=self.headers, timeout=5)
if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
try:
# Create image from response content
from io import BytesIO
logo = Image.open(BytesIO(response.content))
# Verify it's a valid image
logo.verify()
# Reopen after verify
logo = Image.open(BytesIO(response.content))
except Exception as e:
logger.warning(f"Invalid image data from alternative source for {symbol}: {str(e)}")
return None
else:
return None
else: else:
# Try alternative source # Try alternative source
response = requests.get(alt_url, headers=self.headers, timeout=5) response = requests.get(alt_url, headers=self.headers, timeout=5)
if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''): if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
# Create image from response content try:
from io import BytesIO # Create image from response content
logo = Image.open(BytesIO(response.content)) from io import BytesIO
logo = Image.open(BytesIO(response.content))
# Verify it's a valid image
logo.verify()
# Reopen after verify
logo = Image.open(BytesIO(response.content))
except Exception as e:
logger.warning(f"Invalid image data from alternative source for {symbol}: {str(e)}")
return None
else: else:
return None return None
else: else:
# Normal case: open the saved logo file # Normal case: open the saved logo file
logo = Image.open(logo_path) try:
logo = Image.open(logo_path)
# Verify it's a valid image
logo.verify()
# Reopen after verify
logo = Image.open(logo_path)
except Exception as e:
logger.warning(f"Invalid image file for {symbol}: {str(e)}")
return None
# Convert to RGBA if not already # Convert to RGBA if not already
if logo.mode != 'RGBA': if logo.mode != 'RGBA':