updating some examples

Commit 54bdee2 · patx · 2025-01-31T01:56:54-05:00

Changeset
54bdee2a6f11f3633c2b5f85ccfab4cc29d07728
Parents
11bcbae266411a12e2eea905e023302678774db6

View source at this commit

Comments

No comments yet.

Log in to comment

Diff

diff --git a/examples/file_uploads/app.py b/examples/file_uploads/app.py
index e6d24ca..9b13911 100644
--- a/examples/file_uploads/app.py
+++ b/examples/file_uploads/app.py
@@ -1,52 +1,60 @@
 from MicroPie import Server
-import os, uuid, asyncio
 
+import os
+import uuid
+import asyncio
+
+UPLOAD_DIR = "uploads"
+ALLOWED_FILE_TYPES = {"image/png", "image/jpeg", "application/pdf", "video/mp4"}
+MAX_FILE_SIZE = 50 * 1024 * 1024  # 50MB
 
 class Root(Server):
 
     async def save_file_async(self, upload_path, data):
         """
-        Asynchronously writes file data to the specified path.
+        Asynchronously writes file data to the specified path in chunks.
         """
         loop = asyncio.get_running_loop()
         try:
-            await loop.run_in_executor(None, self._write_file, upload_path, data)
+            with open(upload_path, "wb") as f:
+                # Stream the file in chunks (64KB each)
+                chunk_size = 64 * 1024
+                for i in range(0, len(data), chunk_size):
+                    chunk = data[i : i + chunk_size]
+                    await loop.run_in_executor(None, f.write, chunk)
         except IOError as e:
             return 500, f"Failed to save file: {str(e)}"
         return f"File uploaded successfully as '{os.path.basename(upload_path)}'. <a href='/'>Upload another</a>"
 
-    def _write_file(self, upload_path, data):
+    async def upload_file(self):
         """
-        Synchronous function to write data to a file.
-        Used with run_in_executor to avoid blocking the event loop.
+        Handles file upload securely and streams it to disk.
         """
-        with open(upload_path, 'wb') as f:
-            f.write(data)
+        request = self.request  # Access the current request from MicroPie
+        if "file" not in request.files:
+            return 400, "No file uploaded."
 
-    async def upload_file(self, file):
-        """
-        Asynchronous handler function to process uploaded files.
+        file = request.files["file"]
+        filename = os.path.basename(file["filename"])
+        content_type = file["content_type"]
+        data = file["data"]
 
-        Parameters:
-        - file: A dictionary containing 'filename', 'content_type', and 'data'.
-        """
-        if not file:
-            return 400, "No file uploaded."
+        # Security check: Validate file type
+        if content_type not in ALLOWED_FILE_TYPES:
+            return 400, f"Invalid file type: {content_type}. Allowed types: {', '.join(ALLOWED_FILE_TYPES)}"
 
-        filename = file['filename']
-        content_type = file['content_type']
-        data = file['data']
+        # Check file size
+        if len(data) > MAX_FILE_SIZE:
+            return 400, f"File too large! Maximum allowed size is {MAX_FILE_SIZE / (1024 * 1024)} MB."
 
-        # Ensure the 'uploads' directory exists
-        upload_dir = 'uploads'
-        os.makedirs(upload_dir, exist_ok=True)
+        # Ensure upload directory exists
+        os.makedirs(UPLOAD_DIR, exist_ok=True)
 
-        # Sanitize and create a unique filename
-        filename = os.path.basename(filename)
+        # Generate a unique filename
         unique_filename = f"{uuid.uuid4()}_{filename}"
-        upload_path = os.path.join(upload_dir, unique_filename)
+        upload_path = os.path.join(UPLOAD_DIR, unique_filename)
 
-        # Asynchronously save the file
+        # Stream file data to disk in chunks
         return await self.save_file_async(upload_path, data)
 
     def index(self):
@@ -67,6 +75,4 @@ class Root(Server):
             "</html>"
         )
 
-
 app = Root()
-