Update fetch_url_util.py

fetch_url_util.py

@@ -1,4 +1,6 @@
 import io
+import os
+import hashlib
 import urllib.parse
 from PIL import Image
 from curl_cffi import requests
@@ -8,10 +10,26 @@ try:
 except ImportError:
     pass
 
-# Max file size: 150MB
+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
@@ -20,6 +38,23 @@ def fetch_image_from_url(url):
     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}/"
 
@@ -44,25 +79,29 @@ def fetch_image_from_url(url):
         
         response.raise_for_status()
 
-        # 1. Check Content-Length header
         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)")
 
-        # 2. Download in chunks
         downloaded_data = io.BytesIO()
         size = 0
-        for chunk in response.iter_content(chunk_size=65536): # 64KB chunks
+        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)
 
-        # 3. Open and verify
         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:
@@ -76,6 +115,5 @@ def fetch_image_from_url(url):
             raise RuntimeError(error_msg)
             
     finally:
-        # Manually close the response since we used stream=True
         if response is not None:
             response.close()