Skip to content

Commit

Permalink
Add size attribute to UploadFile
Browse files Browse the repository at this point in the history
  • Loading branch information
rafalp committed Feb 4, 2023
1 parent 3697c8d commit 3945635
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 3 deletions.
1 change: 1 addition & 0 deletions docs/requests.md
Expand Up @@ -123,6 +123,7 @@ multidict, containing both file uploads and text input. File upload items are re
* `content_type`: A `str` with the content type (MIME type / media type) (e.g. `image/jpeg`).
* `file`: A <a href="https://docs.python.org/3/library/tempfile.html#tempfile.SpooledTemporaryFile" target="_blank">`SpooledTemporaryFile`</a> (a <a href="https://docs.python.org/3/glossary.html#term-file-like-object" target="_blank">file-like</a> 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`.
* `size`: An `int` with file's size in bytes.

`UploadFile` has the following `async` methods. They all call the corresponding file methods underneath (using the internal `SpooledTemporaryFile`).

Expand Down
4 changes: 4 additions & 0 deletions starlette/datastructures.py
Expand Up @@ -430,12 +430,14 @@ class UploadFile:
def __init__(
self,
file: typing.BinaryIO,
size: int,
*,
filename: typing.Optional[str] = None,
headers: "typing.Optional[Headers]" = None,
) -> None:
self.filename = filename
self.file = file
self.size = size
self.headers = headers or Headers()

@property
Expand All @@ -449,6 +451,8 @@ def _in_memory(self) -> bool:
return not rolled_to_disk

async def write(self, data: bytes) -> None:
self.size += len(data)

if self._in_memory:
self.file.write(data)
else:
Expand Down
8 changes: 5 additions & 3 deletions tests/test_datastructures.py
Expand Up @@ -275,10 +275,12 @@ def test_queryparams():
async def test_upload_file_file_input():
"""Test passing file/stream into the UploadFile constructor"""
stream = io.BytesIO(b"data")
file = UploadFile(filename="file", file=stream)
file = UploadFile(filename="file", file=stream, size=len(stream))
assert file.size == 4
assert await file.read() == b"data"
await file.write(b" and more data!")
assert await file.read() == b""
assert file.size == 19
await file.seek(0)
assert await file.read() == b"data and more data!"

Expand All @@ -292,7 +294,7 @@ async def test_uploadfile_rolling(max_size: int) -> None:
stream: BinaryIO = SpooledTemporaryFile( # type: ignore[assignment]
max_size=max_size
)
file = UploadFile(filename="file", file=stream)
file = UploadFile(filename="file", file=stream, size=len(stream))
assert await file.read() == b""
await file.write(b"data")
assert await file.read() == b""
Expand All @@ -307,7 +309,7 @@ async def test_uploadfile_rolling(max_size: int) -> None:

def test_formdata():
stream = io.BytesIO(b"data")
upload = UploadFile(filename="file", file=stream)
upload = UploadFile(filename="file", file=stream, size=len(stream))
form = FormData([("a", "123"), ("a", "456"), ("b", upload)])
assert "a" in form
assert "A" not in form
Expand Down

0 comments on commit 3945635

Please sign in to comment.