diff --git a/robyn/__init__.py b/robyn/__init__.py index 08dd0631d..a5b343a85 100644 --- a/robyn/__init__.py +++ b/robyn/__init__.py @@ -24,6 +24,7 @@ from robyn.robyn import FunctionInfo, Headers, HttpMethod, Request, Response, WebSocketConnector, get_version from robyn.router import MiddlewareRouter, MiddlewareType, Router, WebSocketRouter from robyn.types import Directory, JsonBody +from robyn.upload import UploadFile from robyn.ws import WebSocketAdapter, WebSocketDisconnect, create_websocket_decorator __version__ = get_version() @@ -810,5 +811,6 @@ def cors_middleware(request): "WebSocketAdapter", "WebSocketDisconnect", "JsonBody", + "UploadFile", "MCPApp", ] diff --git a/robyn/upload.py b/robyn/upload.py new file mode 100644 index 000000000..574c29684 --- /dev/null +++ b/robyn/upload.py @@ -0,0 +1,71 @@ +import io +import mimetypes +from typing import Optional + + +class UploadFile: + """Represents an uploaded file from a multipart form request. + + Provides convenient access to file metadata and content. Wraps the raw + bytes from ``request.files`` with filename, content type, and file-like + read access. + + Usage:: + + @app.post("/upload") + async def upload(request: Request): + raw_files = request.files + for name, data in raw_files.items(): + file = UploadFile(filename=name, file=data) + contents = file.read() + print(f"Got {file.filename} ({file.content_type}): {file.size} bytes") + """ + + __slots__ = ("filename", "content_type", "file", "_size") + + def __init__( + self, + file: bytes, + *, + filename: str = "upload", + content_type: Optional[str] = None, + ) -> None: + self.filename = filename + self.content_type = content_type or self._guess_content_type(filename) + self.file = io.BytesIO(file) if isinstance(file, (bytes, bytearray)) else file + self._size: Optional[int] = len(file) if isinstance(file, (bytes, bytearray)) else None + + @staticmethod + def _guess_content_type(filename: str) -> str: + mime, _ = mimetypes.guess_type(filename) + return mime or "application/octet-stream" + + @property + def size(self) -> int: + if self._size is None: + pos = self.file.tell() + self.file.seek(0, 2) + self._size = self.file.tell() + self.file.seek(pos) + return self._size + + def read(self, size: int = -1) -> bytes: + return self.file.read(size) + + def seek(self, offset: int, whence: int = 0) -> int: + return self.file.seek(offset, whence) + + def tell(self) -> int: + return self.file.tell() + + def close(self) -> None: + self.file.close() + + def __repr__(self) -> str: + return f"UploadFile(filename={self.filename!r}, content_type={self.content_type!r}, size={self.size})" + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() diff --git a/unit_tests/test_upload_file.py b/unit_tests/test_upload_file.py new file mode 100644 index 000000000..b4ee24408 --- /dev/null +++ b/unit_tests/test_upload_file.py @@ -0,0 +1,51 @@ +from robyn.upload import UploadFile + + +def test_upload_file_from_bytes(): + data = b"hello world" + f = UploadFile(data, filename="test.txt") + assert f.filename == "test.txt" + assert f.content_type == "text/plain" + assert f.size == 11 + assert f.read() == b"hello world" + + +def test_upload_file_seek_and_tell(): + data = b"abcdef" + f = UploadFile(data, filename="data.bin") + assert f.tell() == 0 + f.read(3) + assert f.tell() == 3 + f.seek(0) + assert f.read() == b"abcdef" + + +def test_upload_file_content_type_guess(): + f = UploadFile(b"", filename="photo.jpg") + assert f.content_type == "image/jpeg" + + f2 = UploadFile(b"", filename="data.json") + assert f2.content_type == "application/json" + + f3 = UploadFile(b"", filename="unknown") + assert f3.content_type == "application/octet-stream" + + +def test_upload_file_explicit_content_type(): + f = UploadFile(b"", filename="file.txt", content_type="text/csv") + assert f.content_type == "text/csv" + + +def test_upload_file_context_manager(): + data = b"context manager test" + with UploadFile(data, filename="cm.txt") as f: + content = f.read() + assert content == data + + +def test_upload_file_repr(): + f = UploadFile(b"abc", filename="test.txt") + r = repr(f) + assert "test.txt" in r + assert "text/plain" in r + assert "3" in r