Update api.py

api.py

@@ -1,17 +1,19 @@
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import Response, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
+from typing import List, Optional
 import tempfile
 import shutil
 import os
 import subprocess
 import base64
 from pathlib import Path
+import mimetypes
 
 app = FastAPI(
-    title="HTML to PDF API",
-    description="Convert HTML to PDF using Puppeteer",
-    version="1.0.0"
+    title="HTML to PDF API with Image Support",
+    description="Convert HTML to PDF using Puppeteer with image upload support",
+    version="2.0.0"
 )
 
 # Enable CORS
@@ -23,12 +25,49 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-def convert_html_to_pdf(html_content: str, aspect_ratio: str):
+def save_uploaded_images(images: List[UploadFile], temp_dir: str):
+    """Save uploaded images to temp directory and return mapping"""
+    image_mapping = {}
+    images_dir = os.path.join(temp_dir, "images")
+    os.makedirs(images_dir, exist_ok=True)
+    
+    for image in images:
+        if image.filename:
+            # Save image to temp directory
+            image_path = os.path.join(images_dir, image.filename)
+            with open(image_path, 'wb') as f:
+                f.write(image.file.read())
+            
+            # Create mapping for HTML replacement
+            image_mapping[image.filename] = f"images/{image.filename}"
+    
+    return image_mapping
+
+def process_html_with_images(html_content: str, temp_dir: str, image_mapping: dict):
+    """Process HTML to handle image references"""
+    # Replace image references in HTML
+    for original_name, new_path in image_mapping.items():
+        # Handle various image reference patterns
+        patterns = [
+            f'src="{original_name}"',
+            f"src='{original_name}'",
+            f'src={original_name}',
+            f'href="{original_name}"',
+            f"href='{original_name}'"
+        ]
+        
+        for pattern in patterns:
+            if pattern in html_content:
+                html_content = html_content.replace(
+                    pattern,
+                    pattern.replace(original_name, new_path)
+                )
+    
+    return html_content
+
+def convert_html_to_pdf(html_content: str, aspect_ratio: str, temp_dir: str):
     """Convert HTML content to PDF"""
-    temp_dir = None
     try:
-        temp_dir = tempfile.mkdtemp()
-        
         # Style injection for better PDF rendering
         style_injection = """
         <style>
@@ -81,23 +120,21 @@ def convert_html_to_pdf(html_content: str, aspect_ratio: str):
         with open(pdf_file, 'rb') as f:
             pdf_bytes = f.read()
         
-        shutil.rmtree(temp_dir, ignore_errors=True)
         return pdf_bytes
         
     except Exception as e:
-        if temp_dir:
-            shutil.rmtree(temp_dir, ignore_errors=True)
         raise e
 
 @app.get("/")
 async def root():
     """API root endpoint"""
     return {
-        "message": "HTML to PDF Conversion API",
-        "version": "1.0.0",
+        "message": "HTML to PDF Conversion API with Image Support",
+        "version": "2.0.0",
         "endpoints": {
-            "POST /convert": "Convert HTML to PDF (file upload)",
-            "POST /convert-text": "Convert HTML text to PDF",
+            "POST /convert": "Convert HTML to PDF (file upload with optional images)",
+            "POST /convert-text": "Convert HTML text to PDF (with optional image files)",
+            "POST /convert-with-images": "Convert HTML with multiple images",
             "GET /health": "Health check",
             "GET /docs": "API documentation (Swagger UI)"
         }
@@ -111,12 +148,14 @@ async def health_check():
 @app.post("/convert")
 async def convert_file(
     file: UploadFile = File(...),
+    images: Optional[List[UploadFile]] = File(None),
     aspect_ratio: str = Form(default="9:16")
 ):
     """
-    Convert uploaded HTML file to PDF
+    Convert uploaded HTML file to PDF with optional images
     
     - **file**: HTML file to convert
+    - **images**: Optional list of image files (jpg, png, gif, svg, webp)
     - **aspect_ratio**: Page orientation (16:9, 1:1, or 9:16)
     """
     if not file.filename.lower().endswith(('.html', '.htm')):
@@ -125,7 +164,11 @@ async def convert_file(
     if aspect_ratio not in ["16:9", "1:1", "9:16"]:
         raise HTTPException(status_code=400, detail="Invalid aspect ratio. Use: 16:9, 1:1, or 9:16")
     
+    temp_dir = None
     try:
+        # Create temporary directory
+        temp_dir = tempfile.mkdtemp()
+        
         # Read HTML content
         content = await file.read()
         try:
@@ -133,8 +176,16 @@ async def convert_file(
         except UnicodeDecodeError:
             html_content = content.decode('latin-1')
         
+        # Process images if provided
+        if images:
+            image_mapping = save_uploaded_images(images, temp_dir)
+            html_content = process_html_with_images(html_content, temp_dir, image_mapping)
+        
         # Convert to PDF
-        pdf_bytes = convert_html_to_pdf(html_content, aspect_ratio)
+        pdf_bytes = convert_html_to_pdf(html_content, aspect_ratio, temp_dir)
+        
+        # Clean up
+        shutil.rmtree(temp_dir, ignore_errors=True)
         
         # Return PDF file
         filename = file.filename.replace('.html', '.pdf').replace('.htm', '.pdf')
@@ -150,27 +201,43 @@ async def convert_file(
         )
         
     except Exception as e:
+        if temp_dir:
+            shutil.rmtree(temp_dir, ignore_errors=True)
         raise HTTPException(status_code=500, detail=f"Conversion failed: {str(e)}")
 
 @app.post("/convert-text")
 async def convert_text(
     html: str = Form(...),
+    images: Optional[List[UploadFile]] = File(None),
     aspect_ratio: str = Form(default="9:16"),
     return_base64: bool = Form(default=False)
 ):
     """
-    Convert HTML text to PDF
+    Convert HTML text to PDF with optional images
     
     - **html**: HTML content as string
+    - **images**: Optional list of image files
     - **aspect_ratio**: Page orientation (16:9, 1:1, or 9:16)
     - **return_base64**: If true, returns base64 encoded PDF in JSON
     """
     if aspect_ratio not in ["16:9", "1:1", "9:16"]:
         raise HTTPException(status_code=400, detail="Invalid aspect ratio. Use: 16:9, 1:1, or 9:16")
     
+    temp_dir = None
     try:
+        # Create temporary directory
+        temp_dir = tempfile.mkdtemp()
+        
+        # Process images if provided
+        if images:
+            image_mapping = save_uploaded_images(images, temp_dir)
+            html = process_html_with_images(html, temp_dir, image_mapping)
+        
         # Convert to PDF
-        pdf_bytes = convert_html_to_pdf(html, aspect_ratio)
+        pdf_bytes = convert_html_to_pdf(html, aspect_ratio, temp_dir)
+        
+        # Clean up
+        shutil.rmtree(temp_dir, ignore_errors=True)
         
         if return_base64:
             # Return as JSON with base64 encoded PDF
@@ -191,6 +258,78 @@ async def convert_text(
             )
         
     except Exception as e:
+        if temp_dir:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+        raise HTTPException(status_code=500, detail=f"Conversion failed: {str(e)}")
+
+@app.post("/convert-with-images")
+async def convert_with_images(
+    html_file: UploadFile = File(...),
+    images: List[UploadFile] = File(...),
+    aspect_ratio: str = Form(default="9:16")
+):
+    """
+    Convert HTML with multiple images - dedicated endpoint
+    
+    - **html_file**: HTML file to convert
+    - **images**: List of image files (required)
+    - **aspect_ratio**: Page orientation (16:9, 1:1, or 9:16)
+    """
+    if not html_file.filename.lower().endswith(('.html', '.htm')):
+        raise HTTPException(status_code=400, detail="HTML file must be .html or .htm")
+    
+    if aspect_ratio not in ["16:9", "1:1", "9:16"]:
+        raise HTTPException(status_code=400, detail="Invalid aspect ratio. Use: 16:9, 1:1, or 9:16")
+    
+    # Validate image files
+    allowed_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.svg', '.webp', '.bmp'}
+    for img in images:
+        ext = Path(img.filename).suffix.lower()
+        if ext not in allowed_extensions:
+            raise HTTPException(
+                status_code=400, 
+                detail=f"Invalid image format: {img.filename}. Allowed: {', '.join(allowed_extensions)}"
+            )
+    
+    temp_dir = None
+    try:
+        # Create temporary directory
+        temp_dir = tempfile.mkdtemp()
+        
+        # Read HTML content
+        content = await html_file.read()
+        try:
+            html_content = content.decode('utf-8')
+        except UnicodeDecodeError:
+            html_content = content.decode('latin-1')
+        
+        # Save and process images
+        image_mapping = save_uploaded_images(images, temp_dir)
+        html_content = process_html_with_images(html_content, temp_dir, image_mapping)
+        
+        # Convert to PDF
+        pdf_bytes = convert_html_to_pdf(html_content, aspect_ratio, temp_dir)
+        
+        # Clean up
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        
+        # Return PDF
+        filename = html_file.filename.replace('.html', '.pdf').replace('.htm', '.pdf')
+        if not filename.endswith('.pdf'):
+            filename += '.pdf'
+            
+        return Response(
+            content=pdf_bytes,
+            media_type="application/pdf",
+            headers={
+                "Content-Disposition": f"attachment; filename={filename}",
+                "X-Image-Count": str(len(images))
+            }
+        )
+        
+    except Exception as e:
+        if temp_dir:
+            shutil.rmtree(temp_dir, ignore_errors=True)
         raise HTTPException(status_code=500, detail=f"Conversion failed: {str(e)}")
 
 if __name__ == "__main__":