Skip to content

Commit

Permalink
Fix form response file saving
Browse files Browse the repository at this point in the history
  • Loading branch information
CoolCat467 authored and pgjones committed May 1, 2023
1 parent 639d59c commit f9a0281
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 0 deletions.
26 changes: 26 additions & 0 deletions src/quart_trio/datastructures.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
from __future__ import annotations

from os import PathLike

from quart.datastructures import FileStorage
from trio import open_file, Path, wrap_file


class TrioFileStorage(FileStorage):
async def save(self, destination: PathLike, buffer_size: int = 16384) -> None: # type: ignore
wrapped_stream = wrap_file(self.stream)
async with await open_file(destination, "wb") as file_:
data = await wrapped_stream.read(buffer_size)
while data != b"":
await file_.write(data)
data = await wrapped_stream.read(buffer_size)

async def load(self, source: PathLike, buffer_size: int = 16384) -> None:
path = Path(source)
self.filename = path.name
wrapped_stream = wrap_file(self.stream)
async with await open_file(path, "rb") as file_:
data = await file_.read(buffer_size)
while data != b"":
await wrapped_stream.write(data)
data = await file_.read(buffer_size)
7 changes: 7 additions & 0 deletions src/quart_trio/formparser.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
from quart.formparser import FormDataParser

from quart_trio.datastructures import TrioFileStorage


class TrioFormDataParser(FormDataParser):
file_storage_class = TrioFileStorage
3 changes: 3 additions & 0 deletions src/quart_trio/wrappers/request.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from quart.wrappers.request import Body, Request
from werkzeug.exceptions import RequestEntityTooLarge, RequestTimeout

from ..formparser import TrioFormDataParser


class EventWrapper:
def __init__(self) -> None:
Expand Down Expand Up @@ -44,6 +46,7 @@ def __init__(

class TrioRequest(Request):
body_class = TrioBody
form_data_parser_class = TrioFormDataParser
lock_class = trio.Lock # type: ignore

async def get_data(
Expand Down

0 comments on commit f9a0281

Please sign in to comment.