Implement optional chunk encoding in ContentServer

This commit adds an implementation of chunk encoding to ContentServer.
The decision of whether to apply chunking is controlled by an enum,
which takes values YES (apply chunking), NO (don't apply chunking), or
AUTO (apply chunking only if requested by the Transfer-Encoding header).

I set the default value of the flag to NO for backwards compatibility,
but it could be changed to AUTO in the future, which would probably make
more sense.

Author: diazona

pytest_localserver/http.py

@@ -3,6 +3,8 @@
 # This program is release under the MIT license. You can find the full text of
 # the license in the LICENSE file.
 
+import enum
+import itertools
 import json
 import sys
 import threading
@@ -40,6 +42,21 @@ def url(self):
         return '%s://%s:%i' % (proto, host, port)
 
 
+class Chunked(enum.Enum):
+    NO = False
+    YES = True
+    AUTO = None
+
+    def __bool__(self):
+        return bool(self.value)
+
+
+def _encode_chunk(chunk, charset):
+    if isinstance(chunk, str):
+        chunk = chunk.encode(charset)
+    return '{0:x}'.format(len(chunk)).encode(charset) + b'\r\n' + chunk + b'\r\n'
+
+
 class ContentServer(WSGIServer):
 
     """
@@ -68,6 +85,7 @@ def __init__(self, host='127.0.0.1', port=0, ssl_context=None):
         self.show_post_vars = False
         self.compress = None
         self.requests = []
+        self.chunked = Chunked.NO
 
     def __call__(self, environ, start_response):
         """
@@ -81,7 +99,21 @@ def __call__(self, environ, start_response):
         else:
             content = self.content
 
-        response = Response(response=self.content, status=self.code)
+        if (
+            self.chunked == Chunked.YES
+            or (self.chunked == Chunked.AUTO and 'chunked' in self.headers.get('Transfer-encoding', ''))
+        ):
+            # If the code below ever changes to allow setting the charset of
+            # the Response object, the charset used here should also be changed
+            # to match. But until that happens, use UTF-8 since it is Werkzeug's
+            # default.
+            charset = 'utf-8'
+            if isinstance(content, (str, bytes)):
+                content = (_encode_chunk(content, charset), '0\r\n\r\n')
+            else:
+                content = itertools.chain((_encode_chunk(item, charset) for item in content), ['0\r\n\r\n'])
+
+        response = Response(response=content, status=self.code)
         response.headers.clear()
         response.headers.extend(self.headers)
 
@@ -92,14 +124,15 @@ def __call__(self, environ, start_response):
 
         return response(environ, start_response)
 
-    def serve_content(self, content, code=200, headers=None):
+    def serve_content(self, content, code=200, headers=None, chunked=Chunked.NO):
         """
         Serves string content (with specified HTTP error code) as response to
         all subsequent request.
 
         :param content: content to be displayed
         :param code: HTTP status code
         :param headers: HTTP headers to be returned
+        :param chunked: whether to apply chunked transfer encoding to the content
         """
         if not isinstance(content, (str, bytes, list, tuple)):
             # If content is an iterable which is not known to be a string,
@@ -114,6 +147,7 @@ def serve_content(self, content, code=200, headers=None):
                 pass
         self.content = content
         self.code = code
+        self.chunked = chunked
         if headers:
             self.headers = Headers(headers)