Add support for response context manager interface

httpx/_client.py

@@ -1,3 +1,4 @@
+import contextlib
 import datetime
 import enum
 import typing
@@ -8,6 +9,7 @@
 
 from .__version__ import __version__
 from ._auth import Auth, BasicAuth, FunctionAuth
+from ._compat import AsyncExitStack
 from ._config import (
     DEFAULT_LIMITS,
     DEFAULT_MAX_REDIRECTS,
@@ -50,6 +52,8 @@
     NetRCInfo,
     Timer,
     URLPattern,
+    cast_async_context_manager,
+    cast_context_manager,
     get_environment_proxies,
     get_logger,
     same_origin,
@@ -852,19 +856,35 @@ def _send_single_request(self, request: Request, timeout: Timeout) -> Response:
         timer = Timer()
         timer.sync_start()
 
+        exit_stack = contextlib.ExitStack()
+
         with map_exceptions(HTTPCORE_EXC_MAP, request=request):
-            (status_code, headers, stream, ext) = transport.request(
+            # NOTE: Once https://github.com/encode/httpcore/pull/206 is merged,
+            # HTTPCore will always return a context manager.
+            raw_response = transport.request(
                 request.method.encode(),
                 request.url.raw,
                 headers=request.headers.raw,
                 stream=request.stream,  # type: ignore
                 ext={"timeout": timeout.as_dict()},
             )
+            try:
+                (status_code, headers, stream, ext) = raw_response
+            except ValueError:  # pragma: no cover
+                response_ctx = cast_context_manager(raw_response)
+                (status_code, headers, stream, ext) = exit_stack.enter_context(
+                    response_ctx
+                )
+            else:
+                if hasattr(
+                    stream, "close"
+                ):  # Our own content streams don't have close methods.
+                    exit_stack.callback(stream.close)
 
         def on_close(response: Response) -> None:
             response.elapsed = datetime.timedelta(seconds=timer.sync_elapsed())
-            if hasattr(stream, "close"):
-                stream.close()
+            with map_exceptions(HTTPCORE_EXC_MAP, request=request):
+                exit_stack.close()
 
         response = Response(
             status_code,
@@ -1488,20 +1508,38 @@ async def _send_single_request(
         timer = Timer()
         await timer.async_start()
 
+        exit_stack = AsyncExitStack()
+
         with map_exceptions(HTTPCORE_EXC_MAP, request=request):
-            (status_code, headers, stream, ext) = await transport.arequest(
+            # NOTE: Once https://github.com/encode/httpcore/pull/206 is merged,
+            # HTTPCore will always return a context manager.
+            raw_response = transport.arequest(
                 request.method.encode(),
                 request.url.raw,
                 headers=request.headers.raw,
                 stream=request.stream,  # type: ignore
                 ext={"timeout": timeout.as_dict()},
             )
+            try:
+                (status_code, headers, stream, ext) = await raw_response
+            except ValueError:  # pragma: no cover
+                response_ctx = cast_async_context_manager(raw_response)
+                (
+                    status_code,
+                    headers,
+                    stream,
+                    ext,
+                ) = await exit_stack.enter_async_context(response_ctx)
+            else:
+                if hasattr(
+                    stream, "aclose"
+                ):  # Our own content streams don't have close methods.
+                    exit_stack.push_async_callback(stream.aclose)
 
         async def on_close(response: Response) -> None:
             response.elapsed = datetime.timedelta(seconds=await timer.async_elapsed())
-            if hasattr(stream, "aclose"):
-                with map_exceptions(HTTPCORE_EXC_MAP, request=request):
-                    await stream.aclose()
+            with map_exceptions(HTTPCORE_EXC_MAP, request=request):
+                await exit_stack.aclose()
 
         response = Response(
             status_code,

httpx/_compat.py

@@ -0,0 +1,9 @@
+try:
+    from contextlib import AsyncExitStack  # type: ignore  # Py3.6
+except ImportError:  # pragma: no cover
+    # Python 3.6
+    from async_exit_stack import AsyncExitStack  # type: ignore  # noqa: F401
+
+__all__ = [
+    "AsyncExitStack",
+]

httpx/_utils.py

@@ -19,6 +19,8 @@
 if typing.TYPE_CHECKING:  # pragma: no cover
     from ._models import URL
 
+T = typing.TypeVar("T")
+
 
 _HTML5_FORM_ENCODING_REPLACEMENTS = {'"': "%22", "\\": "\\\\"}
 _HTML5_FORM_ENCODING_REPLACEMENTS.update(
@@ -539,3 +541,13 @@ def __eq__(self, other: typing.Any) -> bool:
 
 def warn_deprecated(message: str) -> None:  # pragma: nocover
     warnings.warn(message, DeprecationWarning, stacklevel=2)
+
+
+def cast_context_manager(value: T) -> typing.ContextManager[T]:  # pragma: no cover
+    return value  # type: ignore
+
+
+def cast_async_context_manager(
+    value: typing.Awaitable[T],
+) -> typing.AsyncContextManager[T]:  # pragma: no cover
+    return value  # type: ignore

setup.py

@@ -60,6 +60,8 @@ def get_packages(package):
         "sniffio",
         "rfc3986[idna2008]>=1.3,<2",
         "httpcore>=0.12.1,<0.13",
+        # Backports.
+        "async_exit_stack; python_version<'3.7'",
     ],
     extras_require={
         "http2": "h2==3.*",