From 50d480da1fd11fc606ed1aec75d874ae9d53ea78 Mon Sep 17 00:00:00 2001 From: Elijah Ahianyo Date: Sun, 2 Apr 2023 22:40:05 +0000 Subject: [PATCH] Allow upload event handler to have arbitrary arg name (#755) --- pynecone/.templates/web/utils/state.js | 4 +- pynecone/app.py | 39 ++++++++-- pynecone/components/datadisplay/datatable.py | 1 - pynecone/utils/format.py | 9 +-- tests/conftest.py | 79 ++++++++++++++------ tests/test_app.py | 74 +++++++++++++++++- tests/test_utils.py | 6 +- 7 files changed, 169 insertions(+), 43 deletions(-) diff --git a/pynecone/.templates/web/utils/state.js b/pynecone/.templates/web/utils/state.js index 7cc0532b4..6878bbe6c 100644 --- a/pynecone/.templates/web/utils/state.js +++ b/pynecone/.templates/web/utils/state.js @@ -202,7 +202,6 @@ export const uploadFiles = async ( setResult, files, handler, - multiUpload, endpoint ) => { // If we are already processing an event, or there are no upload files, return. @@ -213,7 +212,6 @@ export const uploadFiles = async ( // Set processing to true to block other events from being processed. setResult({ ...result, processing: true }); - const name = multiUpload ? "files" : "file" const headers = { "Content-Type": files[0].type, }; @@ -221,7 +219,7 @@ export const uploadFiles = async ( // Add the token and handler to the file name. for (let i = 0; i < files.length; i++) { - formdata.append("files", files[i], getToken() + ":" + handler + ":" + name + ":" + files[i].name); + formdata.append("files", files[i], getToken() + ":" + handler + ":" + files[i].name); } // Send the file to the server. diff --git a/pynecone/app.py b/pynecone/app.py index 89e0273ee..8c359e30d 100644 --- a/pynecone/app.py +++ b/pynecone/app.py @@ -1,6 +1,7 @@ """The main Pynecone app.""" -from typing import Any, Callable, Coroutine, Dict, List, Optional, Type, Union +import inspect +from typing import Any, Callable, Coroutine, Dict, List, Optional, Tuple, Type, Union from fastapi import FastAPI, UploadFile from fastapi.middleware import cors @@ -23,7 +24,7 @@ from pynecone.route import ( verify_route_validity, ) from pynecone.state import DefaultState, Delta, State, StateManager, StateUpdate -from pynecone.utils import format +from pynecone.utils import format, types # Define custom types. ComponentCallable = Callable[[], Component] @@ -473,18 +474,46 @@ def upload(app: App): Returns: The state update after processing the event. + + Raises: + ValueError: if there are no args with supported annotation. """ - token, handler, key = files[0].filename.split(":")[:3] + token, handler = files[0].filename.split(":")[:2] for file in files: file.filename = file.filename.split(":")[-1] # Get the state for the session. state = app.state_manager.get_state(token) - # Event payload should have `files` as key for multi-uploads and `file` otherwise + handler_upload_param: Tuple = () + + # get handler function + func = getattr(state, handler.split(".")[-1]) + + # check if there exists any handler args with annotation UploadFile or List[UploadFile] + for k, v in inspect.getfullargspec( + func.fn if isinstance(func, EventHandler) else func + ).annotations.items(): + if ( + types.is_generic_alias(v) + and types._issubclass(v.__args__[0], UploadFile) + or types._issubclass(v, UploadFile) + ): + handler_upload_param = (k, v) + break + + if not handler_upload_param: + raise ValueError( + f"`{handler}` handler should have a parameter annotated with one of the following: List[" + f"pc.UploadFile], pc.UploadFile " + ) + + # check if handler supports multi-upload + multi_upload = types._issubclass(handler_upload_param[1], List) + event = Event( token=token, name=handler, - payload={key: files[0] if key == "file" else files}, + payload={handler_upload_param[0]: files if multi_upload else files[0]}, ) update = await state.process(event) return update diff --git a/pynecone/components/datadisplay/datatable.py b/pynecone/components/datadisplay/datatable.py index 3b4566c67..ac7ede059 100644 --- a/pynecone/components/datadisplay/datatable.py +++ b/pynecone/components/datadisplay/datatable.py @@ -92,7 +92,6 @@ class DataTable(Gridjs): ) def _render(self) -> Tag: - if isinstance(self.data, Var): self.columns = BaseVar( name=f"{self.data.name}.columns" diff --git a/pynecone/utils/format.py b/pynecone/utils/format.py index 1d813463f..f971ab603 100644 --- a/pynecone/utils/format.py +++ b/pynecone/utils/format.py @@ -2,7 +2,6 @@ from __future__ import annotations -import inspect import json import os import re @@ -299,15 +298,9 @@ def format_upload_event(event_spec: EventSpec) -> str: """ from pynecone.compiler import templates - multi_upload = any( - types._issubclass(arg_type, List) - for arg_type in inspect.getfullargspec( - event_spec.handler.fn - ).annotations.values() - ) state, name = get_event_handler_parts(event_spec.handler) parent_state = state.split(".")[0] - return f'uploadFiles({parent_state}, {templates.RESULT}, {templates.SET_RESULT}, {parent_state}.files, "{state}.{name}",{str(multi_upload).lower()} ,UPLOAD)' + return f'uploadFiles({parent_state}, {templates.RESULT}, {templates.SET_RESULT}, {parent_state}.files, "{state}.{name}",UPLOAD)' def format_query_params(router_data: Dict[str, Any]) -> Dict[str, str]: diff --git a/tests/conftest.py b/tests/conftest.py index 44b967e98..fef1876e0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -143,21 +143,13 @@ class UploadState(pc.State): img: str img_list: List[str] - async def handle_upload(self, file: pc.UploadFile): + async def handle_upload1(self, file: pc.UploadFile): """Handle the upload of a file. Args: file: The uploaded file. """ - upload_data = await file.read() - outfile = f".web/public/{file.filename}" - - # Save the file. - with open(outfile, "wb") as file_object: - file_object.write(upload_data) - - # Update the img var. - self.img = file.filename + pass async def multi_handle_upload(self, files: List[pc.UploadFile]): """Handle the upload of a file. @@ -165,16 +157,7 @@ class UploadState(pc.State): Args: files: The uploaded files. """ - for file in files: - upload_data = await file.read() - outfile = f".web/public/{file.filename}" - - # Save the file. - with open(outfile, "wb") as file_object: - file_object.write(upload_data) - - # Update the img var. - self.img_list.append(file.filename) + pass class BaseState(pc.State): @@ -204,7 +187,7 @@ def upload_event_spec(): Returns: Event Spec. """ - return EventSpec(handler=UploadState.handle_upload, upload=True) # type: ignore + return EventSpec(handler=UploadState.handle_upload1, upload=True) # type: ignore @pytest.fixture @@ -225,3 +208,57 @@ def multi_upload_event_spec(): Event Spec. """ return EventSpec(handler=UploadState.multi_handle_upload, upload=True) # type: ignore + + +@pytest.fixture +def upload_state(tmp_path): + """Create upload state. + + Args: + tmp_path: pytest tmp_path + + Returns: + The state + + """ + + class FileUploadState(pc.State): + """The base state for uploading a file.""" + + img: str + img_list: List[str] + + async def handle_upload2(self, file): + """Handle the upload of a file. + + Args: + file: The uploaded file. + """ + upload_data = await file.read() + outfile = f"{tmp_path}/{file.filename}" + + # Save the file. + with open(outfile, "wb") as file_object: + file_object.write(upload_data) + + # Update the img var. + self.img = file.filename + + async def multi_handle_upload(self, files: List[pc.UploadFile]): + """Handle the upload of a file. + + Args: + files: The uploaded files. + """ + for file in files: + upload_data = await file.read() + outfile = f"{tmp_path}/{file.filename}" + + # Save the file. + with open(outfile, "wb") as file_object: + file_object.write(upload_data) + + # Update the img var. + self.img_list.append(file.filename) + + return FileUploadState diff --git a/tests/test_app.py b/tests/test_app.py index eb2a71dc9..62199f209 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -1,13 +1,15 @@ +import io import os.path from typing import List, Tuple, Type import pytest +from fastapi import UploadFile -from pynecone.app import App, DefaultState +from pynecone.app import App, DefaultState, upload from pynecone.components import Box from pynecone.event import Event from pynecone.middleware import HydrateMiddleware -from pynecone.state import State +from pynecone.state import State, StateUpdate from pynecone.style import Style @@ -407,3 +409,71 @@ async def test_dict_mutation_detection__plain_list( ) assert result.delta == expected_delta + + +@pytest.mark.asyncio +async def test_upload_file(upload_state): + """Test that file upload works correctly. + + Args: + upload_state: the state + """ + data = b"This is binary data" + + # Create a binary IO object and write data to it + bio = io.BytesIO() + bio.write(data) + + app = App(state=upload_state) + + file1 = UploadFile( + filename="token:file_upload_state.multi_handle_upload:True:image1.jpg", + file=bio, + content_type="image/jpeg", + ) + file2 = UploadFile( + filename="token:file_upload_state.multi_handle_upload:True:image2.jpg", + file=bio, + content_type="image/jpeg", + ) + fn = upload(app) + result = await fn([file1, file2]) # type: ignore + assert isinstance(result, StateUpdate) + assert result.delta == { + "file_upload_state": {"img_list": ["image1.jpg", "image2.jpg"]} + } + + +@pytest.mark.asyncio +async def test_upload_file_without_annotation(upload_state): + """Test that an error is thrown when there's no param annotated with pc.UploadFile or List[UploadFile]. + + Args: + upload_state: the state + """ + data = b"This is binary data" + + # Create a binary IO object and write data to it + bio = io.BytesIO() + bio.write(data) + + app = App(state=upload_state) + + file1 = UploadFile( + filename="token:upload_state.handle_upload2:True:image1.jpg", + file=bio, + content_type="image/jpeg", + ) + file2 = UploadFile( + filename="token:upload_state.handle_upload2:True:image2.jpg", + file=bio, + content_type="image/jpeg", + ) + fn = upload(app) + with pytest.raises(ValueError) as err: + await fn([file1, file2]) + assert ( + err.value.args[0] + == "`upload_state.handle_upload2` handler should have a parameter annotated with one of the following: " + "List[pc.UploadFile], pc.UploadFile " + ) diff --git a/tests/test_utils.py b/tests/test_utils.py index ea729ea47..f4e1d5762 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -296,7 +296,7 @@ def test_format_upload_event(upload_event_spec): assert ( format.format_upload_event(upload_event_spec) == "uploadFiles(upload_state, result, setResult, " - 'upload_state.files, "upload_state.handle_upload",false ,' + 'upload_state.files, "upload_state.handle_upload1",' "UPLOAD)" ) @@ -310,7 +310,7 @@ def test_format_sub_state_event(upload_sub_state_event_spec): assert ( format.format_upload_event(upload_sub_state_event_spec) == "uploadFiles(base_state, result, setResult, base_state.files, " - '"base_state.sub_upload_state.handle_upload",false ,UPLOAD)' + '"base_state.sub_upload_state.handle_upload",UPLOAD)' ) @@ -323,6 +323,6 @@ def test_format_multi_upload_event(multi_upload_event_spec): assert ( format.format_upload_event(multi_upload_event_spec) == "uploadFiles(upload_state, result, setResult, " - 'upload_state.files, "upload_state.multi_handle_upload",true ,' + 'upload_state.files, "upload_state.multi_handle_upload",' "UPLOAD)" )