Images uploading

kompoth · Apr 2, 2024 · 0be3e5a · 0be3e5a
1 parent bf87df1
commit 0be3e5a
Show file tree

Hide file tree

Showing 8 changed files with 130 additions and 57 deletions.
diff --git a/muckraker/main.py b/muckraker/main.py
@@ -1,12 +1,20 @@
-from fastapi import FastAPI, Response
+import aiofiles
+import asyncio
+from fastapi import FastAPI, Response, UploadFile, File, Form
+from fastapi.exceptions import HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-import tempfile
-from pathlib import Path
 from io import BytesIO
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import List, IO
 
 from .models import Issue
 from .render import render_issue
 
+MAX_IMAGE_SIZE = 2 * 1024 * 1024  # 2 MB
+IMAGE_BATCH = 1024
+ACCEPTED_FILE_TYPES = ("image/png", "image/jpeg", "image/jpg")
+
 app = FastAPI(root_path="/api")
 origins = ["*"]
 app.add_middleware(
@@ -16,15 +24,46 @@
 )
 
 
+async def process_image(image: IO, dir_path: Path):
+    # Validate image
+    if image.content_type not in ACCEPTED_FILE_TYPES:
+        detail = f"Invalid file type: {image.filename}"
+        raise HTTPException(415, detail=detail)
+    if image.size > MAX_IMAGE_SIZE:
+        detail = f"File is too large: {image.filename}"
+        raise HTTPException(413, detail=detail)
+
+    # Save image to the disk
+    image_path = dir_path / image.filename
+    async with aiofiles.open(image_path, "wb") as fd:
+        while content := await image.read(IMAGE_BATCH):
+            await fd.write(content)
+
+
 @app.post("/issue/")
-async def create_issue(issue: Issue):
-    with tempfile.TemporaryDirectory() as tempdir:
-        # Render PDF and save it in the tempdir
-        pdf_tmp_path = Path(tempdir) / "out.pdf"
-        render_issue(issue.config.model_dump(), issue.body, pdf_tmp_path)
+async def create_issue(
+    issue: Issue = Form(),
+    images: List[UploadFile] = []
+):
+    with TemporaryDirectory() as temp_dir:
+        temp_dir_path = Path(temp_dir)
+
+        # Asynchronously process files
+        tasks = [process_image(image, temp_dir_path) for image in images]
+        await asyncio.gather(*tasks)
+
+        # Render PDF and save it in the temp_dir
+        pdf_path = temp_dir_path / "out.pdf"
+        render_issue(
+            config=issue.config.model_dump(),
+            heading=issue.heading.model_dump(),
+            body=issue.body,
+            output=pdf_path,
+            image_dir=temp_dir_path
+        )
 
         # Save PDF to the buffer
-        with open(pdf_tmp_path, "rb") as fd:
+        with open(pdf_path, "rb") as fd:
             buf = BytesIO(fd.read())
 
     # Get it from the buffer

diff --git a/muckraker/md_extensions.py b/muckraker/md_extensions.py
@@ -1,5 +1,4 @@
 import re
-from typing import Any
 from pathlib import Path
 from markdown import Markdown
 from markdown.extensions import Extension

diff --git a/muckraker/render.py b/muckraker/render.py
@@ -43,7 +43,12 @@ def render_issue(
 
     # Render HTML
     issue_template = jinja_env.get_template("newspaper.html")
-    html = issue_template.render(config=config, heading=heading, body=body)
+    html = issue_template.render(
+        config=config,
+        heading=heading,
+        body=body,
+        static="file://" + str(STATIC.resolve())
+    )
 
     # Render PDF
     font_config = FontConfiguration()

diff --git a/muckraker/static/templates/newspaper.html b/muckraker/static/templates/newspaper.html
@@ -2,7 +2,7 @@
 <html>
     <head>
         <meta charset="utf-8"/>
-        <link rel="stylesheet" href="static/style.css" type="text/css" charset="utf-8"/>
+        <link rel="stylesheet" href="{{ static }}/style.css" type="text/css" charset="utf-8"/>
         <style>
             @media print { 
                 @page {

diff --git a/requirements.txt b/requirements.txt
@@ -1,9 +1,11 @@
+aiofiles==23.2.1
 fastapi==0.110.0
 httpx==0.26.0
 Jinja2==3.1.3
 Markdown==3.6
 nh3==0.2.17
 pydantic==2.6.4
 pytest==8.1.1
+python-multipart==0.0.9
 uvicorn==0.29.0
 weasyprint==61.2
diff --git a/static/script.js b/static/script.js
@@ -27,30 +27,33 @@ function updateImageList() {
 
 async function requestPDF() {
     document.getElementById("print-button").disabled = true;
+    var formData = new FormData();
 
-    const requestBody = {
+    const issue = {
         config: {
             size: document.getElementById("size-select").value,
-            bg: document.getElementById("bg-select").value,
-            heading: {
-                title: document.getElementById("title-input").value,
-                subtitle: document.getElementById("subtitle-input").value,
-                no: document.getElementById("issue-no-input").value,
-                date: document.getElementById("issue-date-input").value,
-                cost: document.getElementById("issue-cost-input").value
-            }
+            bg: document.getElementById("bg-select").value
+        },
+        heading: {
+            title: document.getElementById("title-input").value,
+            subtitle: document.getElementById("subtitle-input").value,
+            no: document.getElementById("issue-no-input").value,
+            date: document.getElementById("issue-date-input").value,
+            cost: document.getElementById("issue-cost-input").value
         },
         body: document.getElementById("issue-body-textarea").value
     }
+    formData.append("issue", JSON.stringify(issue));
+
+    const imageInput = document.getElementById("image-input");
+    Array.from(imageInput.files).slice(0, 4).forEach((file) => {
+        formData.append("images", file);
+    });
 
     try {
-        const resp = await fetch("/api/issue", {
+        const resp = await fetch("/api/issue/", {
             method: "POST",
-            headers: {
-                "Content-Type": "application/json",
-                "accept": "application/json"
-            },
-            body: JSON.stringify(requestBody)
+            body: formData 
         });
 
         if (resp.ok) {

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,3 +1,5 @@
+import copy
+from pathlib import Path
 import pytest
 
 LOREM = (
@@ -6,21 +8,34 @@
     "enim ad minim veniam, quis nostrud exercitation ullamco laboris "
     "nisi ut aliquip ex ea commodo consequat."
 )
+__ISSUE_DICT = {
+    "config": {
+        "size": "demitab",
+        "bg": None
+    },
+    "heading": {
+        "title": "Muckraker",
+        "subtitle": "Test sample",
+        "no": "№ 22",
+        "date": "April 1,9999",
+        "cost": "Price 1 c.p."
+    },
+    "body": LOREM
+}
 
 
 @pytest.fixture
-def issue():
-    return {
-        "config": {
-            "size": "demitab",
-            "bg": None,
-            "heading": {
-                "title": "Muckraker",
-                "subtitle": "Test sample",
-                "no": "№ 22",
-                "date": "April 1,9999",
-                "cost": "Price 1 c.p."
-            }
-        },
-        "body": LOREM
-    }
+def issue_dict():
+    return copy.deepcopy(__ISSUE_DICT)
+
+
+@pytest.fixture
+def good_image():
+    path = Path(__file__).parent / "media" / "rufino-train.png"
+    return open(path.resolve(), "rb")
+
+
+@pytest.fixture
+def thick_image():
+    path = Path(__file__).parent / "media" / "nasa-hubble.jpg"
+    return open(path.resolve(), "rb")
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -1,27 +1,37 @@
-import copy
+import json
 from fastapi.testclient import TestClient
 
 from muckraker.main import app
 
+# Is it a good idea to set headers manually?
+FORM_HEADERS = {'Content-Type': 'application/x-www-form-urlencoded'}
 
 client = TestClient(app)
 
 
-def test_issue_thick_body(issue):
-    issue["body"] = issue["body"] * 100
-    resp = client.post("/issue/", json=issue)
-    assert resp.status_code == 422
+def test_issue_correct(issue_dict):
+    content = "issue=" + json.dumps(issue_dict)
+    resp = client.post(
+        "/issue/",
+        content=content,
+        headers=FORM_HEADERS
+    )
+    assert resp.status_code == 200
 
 
-def test_issue_thick_heading(issue):
-    for field in issue["config"]["heading"]:
-        print(field)
-        new_issue = copy.copy(issue)
-        new_issue["config"]["heading"][field] = "a" * 1000
-        resp = client.post("/issue/", json=new_issue)
-        assert resp.status_code == 422
+def test_issue_thick_body(issue_dict):
+    issue_dict["body"] = issue_dict["body"] * 100
+    content = "issue=" + json.dumps(issue_dict)
+    resp = client.post("/issue/", content=content, headers=FORM_HEADERS)
+    assert resp.status_code == 422
 
 
-def test_issue_correct(issue):
-    resp = client.post("/issue/", json=issue)
-    assert resp.status_code == 200
+def test_issue_thick_heading(issue_dict):
+    any_field = list(issue_dict["heading"].keys())[0]
+    issue_dict["heading"][any_field] = "a" * 100
+    content = "issue=" + json.dumps(issue_dict)
+    resp = client.post("/issue/", content=content, headers=FORM_HEADERS)
+    assert resp.status_code == 422
+    detail = resp.json()["detail"][0]
+    assert detail["type"] == "string_too_long"
+    assert detail["loc"] == ["body", "issue", "heading", any_field]