fetch_url_util.py

import io
import os
import hashlib
import urllib.parse
from PIL import Image
from curl_cffi import requests

try:
    import pillow_jxl
except ImportError:
    pass

CACHE_DIR = "/tmp/image_fetch_cache"
MAX_CACHE_COUNT = 100
MAX_SIZE = 150 * 1024 * 1024 
MAX_SIZE_MB = MAX_SIZE / (1024 * 1024)

# Ensure cache directory exists
os.makedirs(CACHE_DIR, exist_ok=True)

def _manage_cache_limit():
    """Removes the oldest files if cache exceeds MAX_CACHE_COUNT."""
    files = [os.path.join(CACHE_DIR, f) for f in os.listdir(CACHE_DIR)]
    if len(files) > MAX_CACHE_COUNT:
        # Sort by access time (oldest first)
        files.sort(key=os.path.getatime)
        for i in range(len(files) - MAX_CACHE_COUNT):
            try:
                os.remove(files[i])
            except OSError:
                pass

def fetch_image_from_url(url):
    if not url:
        return None
    
    url = url.strip()
    if not url.startswith(("http://", "https://")):
        url = "https://" + url

    # Generate a unique filename based on the URL
    url_hash = hashlib.sha256(url.encode()).hexdigest()
    cache_path = os.path.join(CACHE_DIR, f"{url_hash}.png")

    # 1. Check if image is in cache
    if os.path.exists(cache_path):
        try:
            img = Image.open(cache_path)
            img.load()
            # Update access time so it's not cleared by cache manager
            os.utime(cache_path, None)
            return img.convert("RGB")
        except Exception:
            # If cached file is corrupt, remove it and proceed to fetch
            os.remove(cache_path)

    # 2. If not in cache, prepare to fetch
    parsed_url = urllib.parse.urlparse(url)
    base_domain = f"{parsed_url.scheme}://{parsed_url.netloc}/"

    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.7559.59 Safari/537.36",
        "Accept": "image/avif,image/webp,image/jxl,image/apng,image/png,image/jpeg,image/gif,image/svg+xml,image/*,*/*;q=0.8",
        "Accept-Language": "en-US,en;q=0.9",
        "Referer": base_domain,
        "Cache-Control": "no-cache",
    }

    response = None
    try:
        response = requests.get(
            url, 
            headers=headers, 
            impersonate="chrome120", 
            timeout=15,
            allow_redirects=True,
            stream=True 
        )
        
        response.raise_for_status()

        cl = response.headers.get("Content-Length")
        if cl and int(cl) > MAX_SIZE:
            raise RuntimeError(f"File too large (limit: {MAX_SIZE_MB:.0f}MB)")

        downloaded_data = io.BytesIO()
        size = 0
        for chunk in response.iter_content(chunk_size=65536):
            size += len(chunk)
            if size > MAX_SIZE:
                raise RuntimeError(f"File exceeds {MAX_SIZE_MB:.0f}MB limit.")
            downloaded_data.write(chunk)

        downloaded_data.seek(0)
        img = Image.open(downloaded_data)
        img.load()
        
        # 3. Save to cache
        try:
            img.save(cache_path, "PNG")
            _manage_cache_limit()
        except Exception as cache_err:
            print(f"Cache write failed: {cache_err}")

        return img.convert("RGB")
    
    except Exception as e:
        error_msg = f"Fetch Error: {str(e)}"
        print(error_msg)
        
        try:
            import gradio as gr
            raise gr.Error(error_msg)
        except ImportError:
            raise RuntimeError(error_msg)
            
    finally:
        if response is not None:
            response.close()