Skip to content

Commit

Permalink
Fix test suite
Browse files Browse the repository at this point in the history
  • Loading branch information
@Kludex
Kludex committed Sep 22, 2024
1 parent c3ada40 commit 879f91d
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 61 deletions.
74 changes: 14 additions & 60 deletions starlette/responses.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -336,25 +336,15 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:

if len(ranges) == 1:
start, end = ranges[0]
await self._handle_single_range(
send, start, end, stat_result.st_size, send_header_only
)
await self._handle_single_range(send, start, end, stat_result.st_size, send_header_only)
else:
await self._handle_multiple_ranges(
send, ranges, stat_result.st_size, send_header_only
)
await self._handle_multiple_ranges(send, ranges, stat_result.st_size, send_header_only)

if self.background is not None:
await self.background()

async def _handle_simple(self, send: Send, send_header_only: bool) -> None:
await send(
{
"type": "http.response.start",
"status": self.status_code,
"headers": self.raw_headers,
}
)
await send({"type": "http.response.start", "status": self.status_code, "headers": self.raw_headers})
if send_header_only:
await send({"type": "http.response.body", "body": b"", "more_body": False})
else:
Expand All @@ -363,25 +353,13 @@ async def _handle_simple(self, send: Send, send_header_only: bool) -> None:
while more_body:
chunk = await file.read(self.chunk_size)
more_body = len(chunk) == self.chunk_size
await send(
{
"type": "http.response.body",
"body": chunk,
"more_body": more_body,
}
)
await send({"type": "http.response.body", "body": chunk, "more_body": more_body})

async def _handle_single_range(
self, send: Send, start: int, end: int, file_size: int, send_header_only: bool
) -> None:
self.headers["content-range"] = f"bytes {start}-{end - 1}/{file_size}"
await send(
{
"type": "http.response.start",
"status": 206,
"headers": self.raw_headers,
}
)
await send({"type": "http.response.start", "status": 206, "headers": self.raw_headers})
if send_header_only:
await send({"type": "http.response.body", "body": b"", "more_body": False})
else:
Expand All @@ -392,31 +370,19 @@ async def _handle_single_range(
chunk = await file.read(min(self.chunk_size, end - start))
start += len(chunk)
more_body = len(chunk) == self.chunk_size and start < end
await send(
{
"type": "http.response.body",
"body": chunk,
"more_body": more_body,
}
)
await send({"type": "http.response.body", "body": chunk, "more_body": more_body})

async def _handle_multiple_ranges(
self,
send: Send,
ranges: typing.List[typing.Tuple[int, int]],
ranges: list[tuple[int, int]],
file_size: int,
send_header_only: bool,
) -> None:
boundary = md5_hexdigest(os.urandom(16), usedforsecurity=False)
content_type = f"multipart/byteranges; boundary={boundary}"
self.headers["content-type"] = content_type
await send(
{
"type": "http.response.start",
"status": 206,
"headers": self.raw_headers,
}
)
await send({"type": "http.response.start", "status": 206, "headers": self.raw_headers})
if send_header_only:
await send({"type": "http.response.body", "body": b"", "more_body": False})
else:
Expand All @@ -432,18 +398,14 @@ async def _handle_multiple_ranges(
await send(
{
"type": "http.response.body",
"body": f"Content-Type: {self.media_type}\r\n".encode(
"utf-8"
),
"body": f"Content-Type: {self.media_type}\r\n".encode(),
"more_body": True,
}
)
await send(
{
"type": "http.response.body",
"body": f"Content-Range: bytes {start}-{end - 1}/{file_size}\r\n\r\n".encode(
"utf-8"
),
"body": f"Content-Range: bytes {start}-{end - 1}/{file_size}\r\n\r\n".encode(),
"more_body": True,
}
)
Expand All @@ -453,13 +415,7 @@ async def _handle_multiple_ranges(
chunk = await file.read(min(self.chunk_size, end - start))
start += len(chunk)
more_body = len(chunk) == self.chunk_size and start < end
await send(
{
"type": "http.response.body",
"body": chunk,
"more_body": more_body,
}
)
await send({"type": "http.response.body", "body": chunk, "more_body": more_body})
await send(
{
"type": "http.response.body",
Expand All @@ -468,10 +424,8 @@ async def _handle_multiple_ranges(
}
)

def _parse_range_header(
self, http_range: typing.Optional[str], file_size: int
) -> typing.List[typing.Tuple[int, int]]:
ranges: typing.List[typing.Tuple[int, int]] = []
def _parse_range_header(self, http_range: str | None, file_size: int) -> list[tuple[int, int]]:
ranges: list[tuple[int, int]] = []
if http_range is None:
return ranges

Expand Down Expand Up @@ -499,7 +453,7 @@ def _parse_range_header(
return []
ranges.append((start, file_size))
else:
start, end = [int(v) for v in val.split("-", 1)]
start, end = (int(v) for v in val.split("-", 1))
start = int(start)
end = int(end) + 1
if start >= end:
Expand Down
2 changes: 1 addition & 1 deletion tests/test_responses.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,7 @@ async def send(message: Message) -> None:

# Since the TestClient drops the response body on HEAD requests, we need to test
# this directly.
await app({"type": "http", "method": "head"}, receive, send)
await app({"type": "http", "method": "head", "headers": [(b"key", b"value")]}, receive, send)


def test_file_response_set_media_type(tmp_path: Path, test_client_factory: TestClientFactory) -> None:
Expand Down

0 comments on commit 879f91d

Please sign in to comment.