http1: add assemble_body function

3 files changed, 17 insertions, 13 deletions

diff --git a/netlib/exceptions.py b/netlib/exceptions.py
index e30235af..05f1054b 100644
--- a/netlib/exceptions.py
+++ b/netlib/exceptions.py
@@ -40,8 +40,6 @@ class TcpDisconnect(TcpException, Disconnect):
     pass
 
 
-
-
 class TcpReadIncomplete(TcpException):
     pass
 
diff --git a/netlib/http/http1/__init__.py b/netlib/http/http1/__init__.py
index 2d33ff8a..2aa7e26a 100644
--- a/netlib/http/http1/__init__.py
+++ b/netlib/http/http1/__init__.py
@@ -9,6 +9,7 @@ from .read import (
 from .assemble import (
     assemble_request, assemble_request_head,
     assemble_response, assemble_response_head,
+    assemble_body,
 )
 
 
@@ -20,4 +21,5 @@ __all__ = [
     "expected_http_body_size",
     "assemble_request", "assemble_request_head",
     "assemble_response", "assemble_response_head",
+    "assemble_body",
 ]
diff --git a/netlib/http/http1/assemble.py b/netlib/http/http1/assemble.py
index 33b9ef25..7252c446 100644
--- a/netlib/http/http1/assemble.py
+++ b/netlib/http/http1/assemble.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, print_function, division
 
 from ... import utils
+import itertools
 from ...exceptions import HttpException
 from .. import CONTENT_MISSING
 
@@ -25,12 +26,23 @@ def assemble_response(response):
     return head + response.body
 
 
-def assemble_response_head(response, preserve_transfer_encoding=False):
+def assemble_response_head(response):
     first_line = _assemble_response_line(response)
-    headers = _assemble_response_headers(response, preserve_transfer_encoding)
+    headers = _assemble_response_headers(response)
     return b"%s\r\n%s\r\n" % (first_line, headers)
 
 
+def assemble_body(headers, body_chunks):
+    if b"chunked" in headers.get(b"transfer-encoding", b"").lower():
+        for chunk in body_chunks:
+            if chunk:
+                yield b"%x\r\n%s\r\n" % (len(chunk), chunk)
+        yield b"0\r\n\r\n"
+    else:
+        for chunk in body_chunks:
+            yield chunk
+
+
 def _assemble_request_line(request, form=None):
     if form is None:
         form = request.form_out
@@ -87,17 +99,9 @@ def _assemble_response_line(response):
     )
 
 
-def _assemble_response_headers(response, preserve_transfer_encoding=False):
-    # TODO: Remove preserve_transfer_encoding
+def _assemble_response_headers(response):
     headers = response.headers.copy()
     for k in response._headers_to_strip_off:
         headers.pop(k, None)
-    if not preserve_transfer_encoding:
-        headers.pop(b"Transfer-Encoding", None)
-
-        # If body is defined (i.e. not None or CONTENT_MISSING),
-        # we now need to set a content-length header.
-        if response.body or response.body == b"":
-            headers[b"Content-Length"] = str(len(response.body)).encode("ascii")
 
     return bytes(headers)