Allow upload event handler to have arbitrary arg name (#755)

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.

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

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"

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]:

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()
+        pass
-        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
 
     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:
+        pass
-            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)
 
 
 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

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 "
+    )

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)"
     )