Fix support for generator-based WSGI apps

httpx/_dispatch/wsgi.py

@@ -1,4 +1,5 @@
 import io
+import itertools
 import typing
 
 from .._config import TimeoutTypes
@@ -7,6 +8,14 @@
 from .base import SyncDispatcher
 
 
+def _skip_leading_empty_chunks(body: typing.Iterable) -> typing.Iterable:
+    body = iter(body)
+    for chunk in body:
+        if chunk:
+            return itertools.chain([chunk], body)
+    return []
+
+
 class WSGIDispatch(SyncDispatcher):
     """
     A custom SyncDispatcher that handles sending requests directly to an WSGI app.
@@ -88,6 +97,9 @@ def start_response(
             seen_exc_info = exc_info
 
         result = self.app(environ, start_response)
+        # This is needed because the status returned by start_response
+        # shouldn't be used until the first non-empty chunk has been served.
+        result = _skip_leading_empty_chunks(result)
 
         assert seen_status is not None
         assert seen_response_headers is not None

tests/test_wsgi.py

@@ -5,18 +5,20 @@
 import httpx
 
 
-def hello_world(environ, start_response):
-    status = "200 OK"
-    output = b"Hello, World!"
+def application_factory(output):
+    def application(environ, start_response):
+        status = "200 OK"
 
-    response_headers = [
-        ("Content-type", "text/plain"),
-        ("Content-Length", str(len(output))),
-    ]
+        response_headers = [
+            ("Content-type", "text/plain"),
+        ]
 
-    start_response(status, response_headers)
+        start_response(status, response_headers)
 
-    return [output]
+        for item in output:
+            yield item
+
+    return application
 
 
 def echo_body(environ, start_response):
@@ -25,7 +27,6 @@ def echo_body(environ, start_response):
 
     response_headers = [
         ("Content-type", "text/plain"),
-        ("Content-Length", str(len(output))),
     ]
 
     start_response(status, response_headers)
@@ -56,7 +57,6 @@ def raise_exc(environ, start_response):
 
     response_headers = [
         ("Content-type", "text/plain"),
-        ("Content-Length", str(len(output))),
     ]
 
     try:
@@ -69,7 +69,7 @@ def raise_exc(environ, start_response):
 
 
 def test_wsgi():
-    client = httpx.Client(app=hello_world)
+    client = httpx.Client(app=application_factory([b"Hello, World!"]))
     response = client.get("http://www.example.org/")
     assert response.status_code == 200
     assert response.text == "Hello, World!"
@@ -93,3 +93,19 @@ def test_wsgi_exc():
     client = httpx.Client(app=raise_exc)
     with pytest.raises(ValueError):
         client.get("http://www.example.org/")
+
+
+def test_wsgi_generator():
+    output = [b"", b"", b"Some content", b" and more content"]
+    client = httpx.Client(app=application_factory(output))
+    response = client.get("http://www.example.org/")
+    assert response.status_code == 200
+    assert response.text == "Some content and more content"
+
+
+def test_wsgi_generator_empty():
+    output = [b"", b"", b"", b""]
+    client = httpx.Client(app=application_factory(output))
+    response = client.get("http://www.example.org/")
+    assert response.status_code == 200
+    assert response.text == ""