Add sync/async close wrapper util

Author: JordonPhillips

packages/smithy-core/src/smithy_core/aio/types.py

@@ -10,6 +10,7 @@
 from ..exceptions import SmithyException
 from ..interfaces import BytesReader
 from .interfaces import AsyncByteStream, StreamingBlob
+from .utils import close
 
 # The default chunk size for iterating streams.
 _DEFAULT_CHUNK_SIZE = 1024
@@ -115,10 +116,7 @@ def closed(self) -> bool:
     async def close(self) -> None:
         """Closes the stream, as well as the underlying stream where possible."""
         self._closed = True
-        if (close := getattr(self._data, "close", None)) is not None:
-            if asyncio.iscoroutine(result := close()):
-                await result
-
+        await close(self._data)
         self._data = None
 
 
@@ -250,12 +248,9 @@ def closed(self) -> bool:
 
     async def close(self) -> None:
         """Closes the stream, as well as the underlying stream where possible."""
-        if (close := getattr(self._data_source, "close", None)) is not None:
-            if asyncio.iscoroutine(result := close()):
-                await result
-
-        self._data_source = None
         self._buffer.close()
+        await close(self._data_source)
+        self._data_source = None
 
 
 class _AsyncByteStreamIterator:

packages/smithy-core/src/smithy_core/aio/utils.py

@@ -1,18 +1,16 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
-from asyncio import sleep
+from asyncio import sleep, iscoroutine
 from collections.abc import AsyncIterable, Iterable
-from typing import TypeVar
+from typing import Any
 
 from ..exceptions import AsyncBodyException
 from ..interfaces import BytesReader
 from ..interfaces import StreamingBlob as SyncStreamingBlob
 from .interfaces import AsyncByteStream, StreamingBlob
 
-_ListEl = TypeVar("_ListEl")
 
-
-async def async_list(lst: Iterable[_ListEl]) -> AsyncIterable[_ListEl]:
+async def async_list[E](lst: Iterable[E]) -> AsyncIterable[E]:
     """Turn an Iterable into an AsyncIterable."""
     for x in lst:
         await sleep(0)
@@ -53,3 +51,10 @@ def read_streaming_blob(body: StreamingBlob) -> bytes:
             raise AsyncBodyException(
                 f"Expected type {SyncStreamingBlob}, but was {type(body)}"
             )
+
+
+async def close(stream: Any) -> None:
+    """Close a stream, awaiting it if it's async."""
+    if (close := getattr(stream, "close", None)) is not None:
+        if iscoroutine(result := close()):
+            await result

packages/smithy-core/tests/unit/aio/test_utils.py

@@ -0,0 +1,24 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from io import BytesIO
+
+from smithy_core.aio.types import AsyncBytesProvider
+from smithy_core.aio.utils import close
+
+
+async def test_close_sync_closeable() -> None:
+    stream = BytesIO()
+    assert not stream.closed
+    await close(stream)
+    assert stream.closed
+
+
+async def test_close_async_closeable() -> None:
+    stream = AsyncBytesProvider()
+    assert not stream.closed
+    await close(stream)
+    assert stream.closed
+
+
+async def test_close_non_closeable() -> None:
+    await close(b"foo")

packages/smithy-http/src/smithy_http/aio/crt.py

@@ -35,6 +35,7 @@
 
 from smithy_core import interfaces as core_interfaces
 from smithy_core.aio.types import AsyncBytesReader
+from smithy_core.aio.utils import close
 from smithy_core.exceptions import MissingDependencyException
 
 from .. import Field, Fields
@@ -361,7 +362,9 @@ async def _marshal_request(
             # asyncio.create_task we'll start the coroutine without having to
             # explicitly await it.
             crt_body = BufferableByteStream()
-            body = AsyncBytesReader(body)
+
+            if not isinstance(body, AsyncIterable):
+                body = AsyncBytesReader(body)
 
             # Start the read task in the background.
             read_task = asyncio.create_task(self._consume_body_async(body, crt_body))
@@ -380,7 +383,7 @@ async def _marshal_request(
         return crt_request, crt_body
 
     async def _consume_body_async(
-        self, source: AsyncBytesReader, dest: "BufferableByteStream"
+        self, source: AsyncIterable[bytes], dest: "BufferableByteStream"
     ) -> None:
         try:
             async for chunk in source:
@@ -389,7 +392,7 @@ async def _consume_body_async(
             dest.close()
             raise
         finally:
-            await source.close()
+            await close(source)
         dest.end_stream()
 
     def __deepcopy__(self, memo: Any) -> "AWSCRTHTTPClient":