Merge pull request 'feat: Cache images for up to 10 seconds' (#2) from feat/cache-snapshots into main

Reviewed-on: #2

main.py

@@ -2,9 +2,38 @@ import cv2
 import yt_dlp
 import numpy as np
 import json
+import time
+import os
+import base64
 
+# Filesystem-based cache for serverless environment
+CACHE_DIR = "/tmp/snapshot_cache"
+CACHE_DURATION = 10  # seconds
+
+def get_cache_path(url):
+    import hashlib
+    url_hash = hashlib.md5(url.encode('utf-8')).hexdigest()
+    return os.path.join(CACHE_DIR, f"{url_hash}.jpg")
 
 def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=False):
+    if not os.path.exists(CACHE_DIR):
+        os.makedirs(CACHE_DIR, exist_ok=True)
+
+    cache_path = get_cache_path(youtube_url)
+    
+    # Check cache first
+    if os.path.exists(cache_path):
+        mtime = os.path.getmtime(cache_path)
+        if time.time() - mtime < CACHE_DURATION:
+            print(f"Returning cached snapshot for {youtube_url} from {cache_path}")
+            if toBuffer:
+                with open(cache_path, "rb") as f:
+                    return f.read()
+            else:
+                import shutil
+                shutil.copy2(cache_path, output_file)
+                return True
+
     # 1. Configure yt-dlp to get the direct stream URL
     ydl_opts = {
         "format": "best",
@@ -29,41 +58,44 @@ def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=Fa
 
     if not cap.isOpened():
         print("Error: Could not open video stream.")
-        return None if toBuffer else None
+        return None
 
     # Read a single frame
     success, frame = cap.read()
 
+    result = None
     if success:
+        # Save to cache first
+        cv2.imwrite(cache_path, frame)
+        
         if toBuffer:
             # Encode frame as JPEG to memory buffer
             ret, buf = cv2.imencode(".jpg", frame)
             if ret:
                 print("Snapshot captured to buffer.")
-                cap.release()
-                return buf.tobytes()
+                result = buf.tobytes()
             else:
                 print("Error: Could not encode frame to buffer.")
-                cap.release()
-                return None
         else:
             # 3. Save the frame as an image
             cv2.imwrite(output_file, frame)
             print(f"Snapshot saved to {output_file}")
+            result = True
     else:
         print("Error: Could not read frame from stream.")
 
     # Cleanup
     cap.release()
+    return result
 
 
 # Shsf Handler
 def main(args):
     route = args.get("route")
     if route == "snapshot":
-        url = args.get("body", "")
+        url_input = args.get("body", "")
         try:
-            url = json.loads(url)
+            body_json = json.loads(url_input)
         except json.JSONDecodeError:
             print("Error: Invalid JSON input.")
             return {
@@ -74,7 +106,7 @@ def main(args):
                     "Content-Type": "application/json",
                 },
             }
-        url = url.get("url", "")
+        url = body_json.get("url", "")
         if not url:
             print("Error: URL not provided.")
             return {
@@ -86,11 +118,30 @@ def main(args):
                 },
             }
         
+        if not os.path.exists(CACHE_DIR):
+            os.makedirs(CACHE_DIR, exist_ok=True)
+            
+        cache_path = get_cache_path(url)
+        if os.path.exists(cache_path):
+            if time.time() - os.path.getmtime(cache_path) < CACHE_DURATION:
+                print(f"Serving {url} from filesystem cache.")
+                with open(cache_path, "rb") as f:
+                    encoded_content = base64.b64encode(f.read()).decode('utf-8')
+                return {
+                    "_shsf": "v2",
+                    "_code": 200,
+                    "_res": {
+                        "message": "done (cached)",
+                        "buffer": encoded_content
+                    },
+                    "_headers": {
+                        "Content-Type": "application/json",
+                    },
+                }
+
         try:
-            capture_stream_snapshot(url, output_file="/tmp/snapshot.jpg")
-            content = None
-            with open("/tmp/snapshot.jpg", "rb") as f:
-                content = f.read()
+            # Use toBuffer=True to get the bytes directly
+            content = capture_stream_snapshot(url, toBuffer=True)
         except Exception as e:
             print(f"Error: {e}")
             return {
@@ -113,8 +164,7 @@ def main(args):
                 },
             }
         
-        # Convert content to base64 for JSON response
-        import base64
+        # The content returned is already bytes if not from cache
         encoded_content = base64.b64encode(content).decode('utf-8')
         
         return {