From 90b5e2becfaea5b1b7979b40398aa4485b3731f3 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 22 Dec 2021 12:47:25 -0600 Subject: [PATCH 1/4] feat: add headers to UploadFile --- starlette/datastructures.py | 11 +++++++- starlette/formparsers.py | 4 +++ tests/test_formparsers.py | 56 +++++++++++++++++++++++++++++++++++++ 3 files changed, 70 insertions(+), 1 deletion(-) diff --git a/starlette/datastructures.py b/starlette/datastructures.py index 17dc46eb6..23f1a6f3e 100644 --- a/starlette/datastructures.py +++ b/starlette/datastructures.py @@ -415,15 +415,24 @@ class UploadFile: """ spool_max_size = 1024 * 1024 + headers: "Headers" def __init__( - self, filename: str, file: typing.IO = None, content_type: str = "" + self, + filename: str, + file: typing.IO = None, + content_type: str = "", + *, + raw_headers: typing.Optional[typing.List[typing.Tuple[bytes, bytes]]] = None, ) -> None: self.filename = filename self.content_type = content_type if file is None: file = tempfile.SpooledTemporaryFile(max_size=self.spool_max_size) self.file = file + self.headers = Headers( + raw=raw_headers or [(b"Content-Type", content_type.encode("latin-1"))] + ) @property def _in_memory(self) -> bool: diff --git a/starlette/formparsers.py b/starlette/formparsers.py index 1614a9d69..9ded28aba 100644 --- a/starlette/formparsers.py +++ b/starlette/formparsers.py @@ -184,6 +184,7 @@ async def parse(self) -> FormData: file: typing.Optional[UploadFile] = None items: typing.List[typing.Tuple[str, typing.Union[str, UploadFile]]] = [] + item_headers: typing.List[typing.Tuple[bytes, bytes]] = [] # Feed the parser with data from the request. async for chunk in self.stream: @@ -195,6 +196,7 @@ async def parse(self) -> FormData: content_disposition = None content_type = b"" data = b"" + item_headers = [] elif message_type == MultiPartMessage.HEADER_FIELD: header_field += message_bytes elif message_type == MultiPartMessage.HEADER_VALUE: @@ -205,6 +207,7 @@ async def parse(self) -> FormData: content_disposition = header_value elif field == b"content-type": content_type = header_value + item_headers.append((field, header_value)) header_field = b"" header_value = b"" elif message_type == MultiPartMessage.HEADERS_FINISHED: @@ -215,6 +218,7 @@ async def parse(self) -> FormData: file = UploadFile( filename=filename, content_type=content_type.decode("latin-1"), + raw_headers=item_headers, ) else: file = None diff --git a/tests/test_formparsers.py b/tests/test_formparsers.py index 8a1174e1d..384a885dc 100644 --- a/tests/test_formparsers.py +++ b/tests/test_formparsers.py @@ -56,6 +56,26 @@ async def multi_items_app(scope, receive, send): await response(scope, receive, send) +async def app_with_headers(scope, receive, send): + request = Request(scope, receive) + data = await request.form() + output = {} + for key, value in data.items(): + if isinstance(value, UploadFile): + content = await value.read() + output[key] = { + "filename": value.filename, + "content": content.decode(), + "content_type": value.content_type, + "headers": list(value.headers.items()), + } + else: + output[key] = value + await request.close() + response = JSONResponse(output) + await response(scope, receive, send) + + async def app_read_body(scope, receive, send): request = Request(scope, receive) # Read bytes, to force request.stream() to return the already parsed body @@ -137,6 +157,42 @@ def test_multipart_request_multiple_files(tmpdir, test_client_factory): } +def test_multipart_request_multiple_files_with_headers(tmpdir, test_client_factory): + path1 = os.path.join(tmpdir, "test1.txt") + with open(path1, "wb") as file: + file.write(b"") + + path2 = os.path.join(tmpdir, "test2.txt") + with open(path2, "wb") as file: + file.write(b"") + + client = test_client_factory(app_with_headers) + with open(path1, "rb") as f1, open(path2, "rb") as f2: + response = client.post( + "/", + files=[ + ("test1", (None, f1)), + ("test2", ("test2.txt", f2, "text/plain", {"x-custom": "f2"})), + ], + ) + assert response.json() == { + "test1": "", + "test2": { + "filename": "test2.txt", + "content": "", + "content_type": "text/plain", + "headers": [ + [ + "content-disposition", + 'form-data; name="test2"; filename="test2.txt"', + ], + ["content-type", "text/plain"], + ["x-custom", "f2"], + ], + }, + } + + def test_multi_items(tmpdir, test_client_factory): path1 = os.path.join(tmpdir, "test1.txt") with open(path1, "wb") as file: From 49c4d1aac2e523b363a76205d3421e4c9517f392 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 22 Dec 2021 12:50:50 -0600 Subject: [PATCH 2/4] simplify --- starlette/datastructures.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/starlette/datastructures.py b/starlette/datastructures.py index 23f1a6f3e..e4ed257da 100644 --- a/starlette/datastructures.py +++ b/starlette/datastructures.py @@ -430,9 +430,7 @@ def __init__( if file is None: file = tempfile.SpooledTemporaryFile(max_size=self.spool_max_size) self.file = file - self.headers = Headers( - raw=raw_headers or [(b"Content-Type", content_type.encode("latin-1"))] - ) + self.headers = Headers(raw=raw_headers or []) @property def _in_memory(self) -> bool: From 79772a0d70af580b8ce7da6edf50c3b7702c5059 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 5 Jan 2022 08:47:23 -0800 Subject: [PATCH 3/4] use headers: Headers in UploadFile.__init__ --- starlette/datastructures.py | 4 ++-- starlette/formparsers.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/starlette/datastructures.py b/starlette/datastructures.py index e4ed257da..64f964a91 100644 --- a/starlette/datastructures.py +++ b/starlette/datastructures.py @@ -423,14 +423,14 @@ def __init__( file: typing.IO = None, content_type: str = "", *, - raw_headers: typing.Optional[typing.List[typing.Tuple[bytes, bytes]]] = None, + headers: "typing.Optional[Headers]" = None, ) -> None: self.filename = filename self.content_type = content_type if file is None: file = tempfile.SpooledTemporaryFile(max_size=self.spool_max_size) self.file = file - self.headers = Headers(raw=raw_headers or []) + self.headers = headers or Headers() @property def _in_memory(self) -> bool: diff --git a/starlette/formparsers.py b/starlette/formparsers.py index 9ded28aba..decaf0bfd 100644 --- a/starlette/formparsers.py +++ b/starlette/formparsers.py @@ -218,7 +218,7 @@ async def parse(self) -> FormData: file = UploadFile( filename=filename, content_type=content_type.decode("latin-1"), - raw_headers=item_headers, + headers=Headers(raw=item_headers), ) else: file = None From baa8216f4340d2322a1381e2855b9b2a8aed4b24 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 5 Jan 2022 08:53:52 -0800 Subject: [PATCH 4/4] Add docs --- docs/requests.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/requests.md b/docs/requests.md index a72cb75dc..f4d867ab1 100644 --- a/docs/requests.md +++ b/docs/requests.md @@ -122,7 +122,7 @@ multidict, containing both file uploads and text input. File upload items are re * `filename`: A `str` with the original file name that was uploaded (e.g. `myimage.jpg`). * `content_type`: A `str` with the content type (MIME type / media type) (e.g. `image/jpeg`). * `file`: A `SpooledTemporaryFile` (a file-like object). This is the actual Python file that you can pass directly to other functions or libraries that expect a "file-like" object. - +* `headers`: A `Headers` object. Often this will only be the `Content-Type` header, but if additional headers were included in the multipart field they will be included here. Note that these headers have no relationship with the headers in `Request.headers`. `UploadFile` has the following `async` methods. They all call the corresponding file methods underneath (using the internal `SpooledTemporaryFile`). @@ -142,6 +142,7 @@ filename = form["upload_file"].filename contents = await form["upload_file"].read() ``` + #### Application The originating Starlette application can be accessed via `request.app`.