http1: add assemble_body function

4 files changed, 31 insertions, 17 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)
diff --git a/test/http/http1/test_assemble.py b/test/http/http1/test_assemble.py
index 8a0a54f1..cdc8bda9 100644
--- a/test/http/http1/test_assemble.py
+++ b/test/http/http1/test_assemble.py
@@ -4,8 +4,8 @@ from netlib.http import CONTENT_MISSING, Headers
 from netlib.http.http1.assemble import (
     assemble_request, assemble_request_head, assemble_response,
     assemble_response_head, _assemble_request_line, _assemble_request_headers,
-    _assemble_response_headers
-)
+    _assemble_response_headers,
+    assemble_body)
 from netlib.tutils import treq, raises, tresp
 
 
@@ -50,6 +50,17 @@ def test_assemble_response_head():
     assert b"message" not in c
 
 
+def test_assemble_body():
+    c = list(assemble_body(Headers(), [b"body"]))
+    assert c == [b"body"]
+
+    c = list(assemble_body(Headers(transfer_encoding="chunked"), [b"123456789a", b""]))
+    assert c == [b"a\r\n123456789a\r\n", b"0\r\n\r\n"]
+
+    c = list(assemble_body(Headers(transfer_encoding="chunked"), [b"123456789a"]))
+    assert c == [b"a\r\n123456789a\r\n", b"0\r\n\r\n"]
+
+
 def test_assemble_request_line():
     assert _assemble_request_line(treq()) == b"GET /path HTTP/1.1"
 
@@ -83,8 +94,7 @@ def test_assemble_response_headers():
     r = tresp(body=b"")
     r.headers["Transfer-Encoding"] = b"chunked"
     c = _assemble_response_headers(r)
-    assert b"Content-Length" in c
-    assert b"Transfer-Encoding" not in c
+    assert b"Transfer-Encoding" in c
 
     assert b"Proxy-Connection" not in _assemble_response_headers(
         tresp(headers=Headers(Proxy_Connection=b"42"))