Allow to disable Transfer-Encoding: chunked

This commit is contained in:
andreymal 2019-04-20 22:26:30 +03:00
parent 6a4a3f617f
commit 9f07109616
No known key found for this signature in database
GPG Key ID: CB605DA977D27A93

View File

@ -59,16 +59,23 @@ class StreamingHTTPResponse(BaseHTTPResponse):
"status", "status",
"content_type", "content_type",
"headers", "headers",
"chunked",
"_cookies", "_cookies",
) )
def __init__( def __init__(
self, streaming_fn, status=200, headers=None, content_type="text/plain" self,
streaming_fn,
status=200,
headers=None,
content_type="text/plain",
chunked=None,
): ):
self.content_type = content_type self.content_type = content_type
self.streaming_fn = streaming_fn self.streaming_fn = streaming_fn
self.status = status self.status = status
self.headers = CIMultiDict(headers or {}) self.headers = CIMultiDict(headers or {})
self.chunked = chunked
self._cookies = None self._cookies = None
async def write(self, data): async def write(self, data):
@ -79,7 +86,10 @@ class StreamingHTTPResponse(BaseHTTPResponse):
if type(data) != bytes: if type(data) != bytes:
data = self._encode_body(data) data = self._encode_body(data)
if self.chunked:
self.protocol.push_data(b"%x\r\n%b\r\n" % (len(data), data)) self.protocol.push_data(b"%x\r\n%b\r\n" % (len(data), data))
else:
self.protocol.push_data(data)
await self.protocol.drain() await self.protocol.drain()
async def stream( async def stream(
@ -88,6 +98,8 @@ class StreamingHTTPResponse(BaseHTTPResponse):
"""Streams headers, runs the `streaming_fn` callback that writes """Streams headers, runs the `streaming_fn` callback that writes
content to the response body, then finalizes the response body. content to the response body, then finalizes the response body.
""" """
if self.chunked is None:
self.chunked = version != "1.0"
headers = self.get_headers( headers = self.get_headers(
version, version,
keep_alive=keep_alive, keep_alive=keep_alive,
@ -96,6 +108,7 @@ class StreamingHTTPResponse(BaseHTTPResponse):
self.protocol.push_data(headers) self.protocol.push_data(headers)
await self.protocol.drain() await self.protocol.drain()
await self.streaming_fn(self) await self.streaming_fn(self)
if self.chunked:
self.protocol.push_data(b"0\r\n\r\n") self.protocol.push_data(b"0\r\n\r\n")
# no need to await drain here after this write, because it is the # no need to await drain here after this write, because it is the
# very last thing we write and nothing needs to wait for it. # very last thing we write and nothing needs to wait for it.
@ -109,6 +122,11 @@ class StreamingHTTPResponse(BaseHTTPResponse):
if keep_alive and keep_alive_timeout is not None: if keep_alive and keep_alive_timeout is not None:
timeout_header = b"Keep-Alive: %d\r\n" % keep_alive_timeout timeout_header = b"Keep-Alive: %d\r\n" % keep_alive_timeout
chunked = self.chunked
if chunked is None:
chunked = version != "1.0"
if chunked:
self.headers["Transfer-Encoding"] = "chunked" self.headers["Transfer-Encoding"] = "chunked"
self.headers.pop("Content-Length", None) self.headers.pop("Content-Length", None)
self.headers["Content-Type"] = self.headers.get( self.headers["Content-Type"] = self.headers.get(
@ -327,6 +345,7 @@ async def file_stream(
mime_type=None, mime_type=None,
headers=None, headers=None,
filename=None, filename=None,
chunked=None,
_range=None, _range=None,
): ):
"""Return a streaming response object with file data. """Return a streaming response object with file data.
@ -336,6 +355,7 @@ async def file_stream(
:param mime_type: Specific mime_type. :param mime_type: Specific mime_type.
:param headers: Custom Headers. :param headers: Custom Headers.
:param filename: Override filename. :param filename: Override filename.
:param chunked: Enable or disable chunked transfer-encoding (default: auto)
:param _range: :param _range:
""" """
headers = headers or {} headers = headers or {}
@ -383,6 +403,7 @@ async def file_stream(
status=status, status=status,
headers=headers, headers=headers,
content_type=mime_type, content_type=mime_type,
chunked=chunked,
) )
@ -391,6 +412,7 @@ def stream(
status=200, status=200,
headers=None, headers=None,
content_type="text/plain; charset=utf-8", content_type="text/plain; charset=utf-8",
chunked=None,
): ):
"""Accepts an coroutine `streaming_fn` which can be used to """Accepts an coroutine `streaming_fn` which can be used to
write chunks to a streaming response. Returns a `StreamingHTTPResponse`. write chunks to a streaming response. Returns a `StreamingHTTPResponse`.
@ -409,9 +431,14 @@ def stream(
writes content to that response. writes content to that response.
:param mime_type: Specific mime_type. :param mime_type: Specific mime_type.
:param headers: Custom Headers. :param headers: Custom Headers.
:param chunked: Enable or disable chunked transfer-encoding (default: auto)
""" """
return StreamingHTTPResponse( return StreamingHTTPResponse(
streaming_fn, headers=headers, content_type=content_type, status=status streaming_fn,
headers=headers,
content_type=content_type,
status=status,
chunked=chunked,
) )