Added .close() to memory streams (#313)

Author: agronholm

docs/streams.rst

@@ -67,6 +67,14 @@ Example::
 
     run(main)
 
+In contrast to other AnyIO streams (but in line with trio's Channels), memory object streams can be
+closed synchronously, using either the ``close()`` method or by using the stream as a context
+manager::
+
+    def synchronous_callback(send_stream: MemoryObjectSendStream) -> None:
+        with stream:
+            stream.send_nowait('hello')
+
 Stapled streams
 ---------------
 

docs/versionhistory.rst

@@ -5,6 +5,8 @@ This library adheres to `Semantic Versioning 2.0 <http://semver.org/>`_.
 
 **UNRELEASED**
 
+- Added the ability to close memory object streams synchronously (including support for use as a
+  synchronous context manager)
 - Fixed ``to_thread.run_sync()`` hanging on the second call on asyncio when used with
   ``loop.run_until_complete()``
 - Fixed the type annotation of ``open_signal_receiver()`` as a synchronous context manager

src/anyio/streams/memory.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict, deque
 from dataclasses import dataclass, field
-from typing import Deque, Generic, List, NamedTuple, TypeVar
+from types import TracebackType
+from typing import Deque, Generic, List, NamedTuple, Optional, Type, TypeVar
 
 from .. import (
     BrokenResourceError, ClosedResourceError, EndOfStream, WouldBlock, get_cancelled_exc_class)
@@ -112,7 +113,14 @@ def clone(self) -> 'MemoryObjectReceiveStream':
 
         return MemoryObjectReceiveStream(_state=self._state)
 
-    async def aclose(self) -> None:
+    def close(self) -> None:
+        """
+        Close the stream.
+
+        This works the exact same way as :meth:`aclose`, but is provided as a special case for the
+        benefit of synchronous callbacks.
+
+        """
         if not self._closed:
             self._closed = True
             self._state.open_receive_channels -= 1
@@ -121,6 +129,9 @@ async def aclose(self) -> None:
                 for event in send_events:
                     event.set()
 
+    async def aclose(self) -> None:
+        self.close()
+
     def statistics(self) -> MemoryObjectStreamStatistics:
         """
         Return statistics about the current state of this stream.
@@ -129,6 +140,14 @@ def statistics(self) -> MemoryObjectStreamStatistics:
         """
         return self._state.statistics()
 
+    def __enter__(self) -> 'MemoryObjectReceiveStream[T_Item]':
+        return self
+
+    def __exit__(self, exc_type: Optional[Type[BaseException]],
+                 exc_val: Optional[BaseException],
+                 exc_tb: Optional[TracebackType]) -> None:
+        self.close()
+
 
 @dataclass(eq=False)
 class MemoryObjectSendStream(Generic[T_Item], ObjectSendStream[T_Item]):
@@ -198,7 +217,14 @@ def clone(self) -> 'MemoryObjectSendStream':
 
         return MemoryObjectSendStream(_state=self._state)
 
-    async def aclose(self) -> None:
+    def close(self) -> None:
+        """
+        Close the stream.
+
+        This works the exact same way as :meth:`aclose`, but is provided as a special case for the
+        benefit of synchronous callbacks.
+
+        """
         if not self._closed:
             self._closed = True
             self._state.open_send_channels -= 1
@@ -208,10 +234,21 @@ async def aclose(self) -> None:
                 for event in receive_events:
                     event.set()
 
+    async def aclose(self) -> None:
+        self.close()
+
     def statistics(self) -> MemoryObjectStreamStatistics:
         """
         Return statistics about the current state of this stream.
 
         .. versionadded:: 3.0
         """
         return self._state.statistics()
+
+    def __enter__(self) -> 'MemoryObjectSendStream[T_Item]':
+        return self
+
+    def __exit__(self, exc_type: Optional[Type[BaseException]],
+                 exc_val: Optional[BaseException],
+                 exc_tb: Optional[TracebackType]) -> None:
+        self.close()

tests/streams/test_memory.py

@@ -320,3 +320,15 @@ async def test_statistics() -> None:
         assert stream.statistics().current_buffer_used == 0
         assert stream.statistics().tasks_waiting_send == 0
         assert stream.statistics().tasks_waiting_receive == 0
+
+
+async def test_sync_close() -> None:
+    send_stream, receive_stream = create_memory_object_stream(1)
+    with send_stream, receive_stream:
+        pass
+
+    with pytest.raises(ClosedResourceError):
+        send_stream.send_nowait(None)
+
+    with pytest.raises(ClosedResourceError):
+        receive_stream.receive_nowait()