Convert FlowControlMixin to StdoutWriterProtocol

Add support for waiting for stream close.
And actually use the new class.

Author: marmarek

splitgpg2/__init__.py

@@ -46,6 +46,8 @@
 from typing import Optional, Dict, Callable, Awaitable, Tuple, Pattern, List, \
      Union, Any, TypeVar, Set, TYPE_CHECKING, Coroutine, Sequence, cast
 
+from .stdiostream import StdoutWriterProtocol
+
 if TYPE_CHECKING:
     from typing_extensions import Protocol
     from typing import TypeAlias
@@ -1405,7 +1407,7 @@ def open_stdinout_connection(*,
 
     write_transport, write_protocol = loop.run_until_complete(
             loop.connect_write_pipe(
-                lambda: asyncio.streams.FlowControlMixin(loop),
+                lambda: StdoutWriterProtocol(loop),
                 sys.stdout.buffer))
     writer = asyncio.StreamWriter(write_transport, write_protocol, None, loop)
 

splitgpg2/stdiostream.py

@@ -2,19 +2,18 @@
 # based on asyncio library:
 # Copyright (C) 2001 Python Software Foundation
 #
+# Copyright (C) 2024 Marek Marczykowski-Górecki
+#                               <marmarek@invisiblethingslab.com>
 #
 
 import collections
 from asyncio import protocols, events
 
-
-class FlowControlMixin(protocols.Protocol):
+class StdoutWriterProtocol(protocols.Protocol):
     """Reusable flow control logic for StreamWriter.drain().
-
     This implements the protocol methods pause_writing(),
     resume_writing() and connection_lost().  If the subclass overrides
     these it must call the super methods.
-
     StreamWriter.drain() must wait for _drain_helper() coroutine.
     """
 
@@ -26,18 +25,15 @@ def __init__(self, loop=None):
         self._paused = False
         self._drain_waiters = collections.deque()
         self._connection_lost = False
+        self._closed = self._loop.create_future()
 
     def pause_writing(self):
         assert not self._paused
         self._paused = True
-        if self._loop.get_debug():
-            logger.debug("%r pauses writing", self)
 
     def resume_writing(self):
         assert self._paused
         self._paused = False
-        if self._loop.get_debug():
-            logger.debug("%r resumes writing", self)
 
         for waiter in self._drain_waiters:
             if not waiter.done():
@@ -55,6 +51,11 @@ def connection_lost(self, exc):
                     waiter.set_result(None)
                 else:
                     waiter.set_exception(exc)
+        if not self._closed.done():
+            if exc is None:
+                self._closed.set_result(None)
+            else:
+                self._closed.set_exception(exc)
 
     async def _drain_helper(self):
         if self._connection_lost:
@@ -68,6 +69,6 @@ async def _drain_helper(self):
         finally:
             self._drain_waiters.remove(waiter)
 
+    # pylint: disable=unused-argument
     def _get_close_waiter(self, stream):
-        raise NotImplementedError
-
+        return self._closed