Allow connection to be used as an async context manager

docs/index.rst

@@ -29,3 +29,16 @@ Additionally, :py:meth:`pywreck.request` is provided for custom HTTP methods.
     :hidden:
 
     api
+
+Reusing Connections
+===================
+
+Creating TCP connections and negotiating TLS can sometimes be expensive. An API
+is provided to allow for multiple HTTP requests on the same connection.
+
+.. code-block:: python
+
+    async with await pywreck.Connection.create("www.example.com") as conn:
+        for _ in range(2):
+            response = await conn.request("HEAD", "/")
+            print(response)

src/pywreck/__init__.py

@@ -16,7 +16,8 @@
 import os.path
 import ssl
 from dataclasses import dataclass
-from typing import Dict, Optional, Union
+from types import TracebackType
+from typing import Dict, Optional, Type, Union
 
 try:
     with open(os.path.join(os.path.dirname(__file__), "version.txt")) as f:
@@ -121,6 +122,9 @@ async def create(
         reader, writer = await asyncio.open_connection(host, port, ssl=ssl)
         return Connection(reader, writer, host, timeout)
 
+    async def __aenter__(self) -> "Connection":
+        return self
+
     async def request(
         self,
         method: str,
@@ -219,6 +223,14 @@ async def request(
 
         return Response(status, response_headers, response_data)
 
+    async def __aexit__(
+        self,
+        exc: Optional[Type[BaseException]],
+        value: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        await self.close()
+
     async def close(self) -> None:
         """Close the connection"""
         writer = self._writer

tests/test_pywreck.py

@@ -123,15 +123,14 @@ async def _async():
             timeout=None,
         )
 
-        for _ in range(2):
-            response = await connection.request(
-                "GET",
-                "/",
-                headers={"user-agent": "pywreck test, yo!"},
-            )
-            validate_response(response)
-
-        await connection.close()
+        async with connection:
+            for _ in range(2):
+                response = await connection.request(
+                    "GET",
+                    "/",
+                    headers={"user-agent": "pywreck test, yo!"},
+                )
+                validate_response(response)
 
     loop.run_until_complete(_async())