Turn create_pipe_input into a context manager (breaking change).

This makes context managers of the following:
- `create_pipe_input()`
- `PosixPipeInput`
- `Win32PipeInput`

The reason for this change is that the close method of the pipe should only
close the write-end, and as a consequence of that, the read-end should trigger
an `EOFError`. Before this change, the read-end was also closed, and that
caused the key input to never wake up and "read" the end-of-file. However, we
still want to close the read end at some point, and that's why this is a
context manager now.

As part of this change, exceptions that are raised in the TelnetServer interact
method, won't cause cause the whole server to crash.

See also: #1585

Co-Author: Frank Wu <kwyd@163.com>

docs/pages/advanced_topics/unit_testing.rst

@@ -49,18 +49,14 @@ In the following example we use a
     from prompt_toolkit.output import DummyOutput
 
     def test_prompt_session():
-        inp = create_pipe_input()
-
-        try:
+        with create_pipe_input() as inp:
             inp.send_text("hello\n")
             session = PromptSession(
                 input=inp,
                 output=DummyOutput(),
             )
 
             result = session.prompt()
-        finally:
-            inp.close()
 
         assert result == "hello"
 
@@ -116,12 +112,9 @@ single fixture that does it for every test. Something like this:
 
     @pytest.fixture(autouse=True, scope="function")
     def mock_input():
-        pipe_input = create_pipe_input()
-        try:
+        with create_pipe_input() as pipe_input:
             with create_app_session(input=pipe_input, output=DummyOutput()):
                 yield pipe_input
-        finally:
-            pipe_input.close()
 
 
 Type checking

prompt_toolkit/contrib/ssh/server.py

@@ -10,7 +10,7 @@
 from prompt_toolkit.application.current import AppSession, create_app_session
 from prompt_toolkit.data_structures import Size
 from prompt_toolkit.eventloop import get_event_loop
-from prompt_toolkit.input import create_pipe_input
+from prompt_toolkit.input import PipeInput, create_pipe_input
 from prompt_toolkit.output.vt100 import Vt100_Output
 
 __all__ = ["PromptToolkitSSHSession", "PromptToolkitSSHServer"]
@@ -28,7 +28,7 @@ def __init__(
         # PipInput object, for sending input in the CLI.
         # (This is something that we can use in the prompt_toolkit event loop,
         # but still write date in manually.)
-        self._input = create_pipe_input()
+        self._input: Optional[PipeInput] = None
         self._output: Optional[Vt100_Output] = None
 
         # Output object. Don't render to the real stdout, but write everything
@@ -88,16 +88,17 @@ async def _interact(self) -> None:
         self._output = Vt100_Output(
             self.stdout, self._get_size, term=term, write_binary=False
         )
-        with create_app_session(input=self._input, output=self._output) as session:
-            self.app_session = session
-            try:
-                await self.interact(self)
-            except BaseException:
-                traceback.print_exc()
-            finally:
-                # Close the connection.
-                self._chan.close()
-                self._input.close()
+        with create_pipe_input() as self._input:
+            with create_app_session(input=self._input, output=self._output) as session:
+                self.app_session = session
+                try:
+                    await self.interact(self)
+                except BaseException:
+                    traceback.print_exc()
+                finally:
+                    # Close the connection.
+                    self._chan.close()
+                    self._input.close()
 
     def terminal_size_changed(
         self, width: int, height: int, pixwidth: object, pixheight: object
@@ -107,6 +108,10 @@ def terminal_size_changed(
             self.app_session.app._on_resize()
 
     def data_received(self, data: str, datatype: object) -> None:
+        if self._input is None:
+            # Should not happen.
+            return
+
         self._input.send_text(data)
 
 

prompt_toolkit/contrib/telnet/server.py

@@ -11,7 +11,7 @@
 from prompt_toolkit.data_structures import Size
 from prompt_toolkit.eventloop import get_event_loop
 from prompt_toolkit.formatted_text import AnyFormattedText, to_formatted_text
-from prompt_toolkit.input import create_pipe_input
+from prompt_toolkit.input import PipeInput, create_pipe_input
 from prompt_toolkit.output.vt100 import Vt100_Output
 from prompt_toolkit.renderer import print_formatted_text as print_formatted_text
 from prompt_toolkit.styles import BaseStyle, DummyStyle
@@ -87,6 +87,7 @@ def __init__(self, connection: socket.socket, encoding: str) -> None:
         self._connection = connection
         self._errors = "strict"
         self._buffer: List[bytes] = []
+        self._closed = False
 
     def write(self, data: str) -> None:
         data = data.replace("\n", "\r\n")
@@ -104,6 +105,9 @@ def flush(self) -> None:
 
         self._buffer = []
 
+    def close(self) -> None:
+        self._closed = True
+
     @property
     def encoding(self) -> str:
         return self._encoding
@@ -126,6 +130,7 @@ def __init__(
         server: "TelnetServer",
         encoding: str,
         style: Optional[BaseStyle],
+        vt100_input: PipeInput,
     ) -> None:
 
         self.conn = conn
@@ -136,6 +141,7 @@ def __init__(
         self.style = style
         self._closed = False
         self._ready = asyncio.Event()
+        self.vt100_input = vt100_input
         self.vt100_output = None
 
         # Create "Output" object.
@@ -144,9 +150,6 @@ def __init__(
         # Initialize.
         _initialize_telnet(conn)
 
-        # Create input.
-        self.vt100_input = create_pipe_input()
-
         # Create output.
         def get_size() -> Size:
             return self.size
@@ -197,12 +200,6 @@ def handle_incoming_data() -> None:
             with create_app_session(input=self.vt100_input, output=self.vt100_output):
                 self.context = contextvars.copy_context()
                 await self.interact(self)
-        except Exception as e:
-            print("Got %s" % type(e).__name__, e)
-            import traceback
-
-            traceback.print_exc()
-            raise
         finally:
             self.close()
 
@@ -222,6 +219,7 @@ def close(self) -> None:
             self.vt100_input.close()
             get_event_loop().remove_reader(self.conn)
             self.conn.close()
+            self.stdout.close()
 
     def send(self, formatted_text: AnyFormattedText) -> None:
         """
@@ -336,22 +334,34 @@ def _accept(self) -> None:
         conn, addr = self._listen_socket.accept()
         logger.info("New connection %r %r", *addr)
 
-        connection = TelnetConnection(
-            conn, addr, self.interact, self, encoding=self.encoding, style=self.style
-        )
-        self.connections.add(connection)
-
         # Run application for this connection.
         async def run() -> None:
-            logger.info("Starting interaction %r %r", *addr)
             try:
-                await connection.run_application()
-            except Exception as e:
-                print(e)
+                with create_pipe_input() as vt100_input:
+                    connection = TelnetConnection(
+                        conn,
+                        addr,
+                        self.interact,
+                        self,
+                        encoding=self.encoding,
+                        style=self.style,
+                        vt100_input=vt100_input,
+                    )
+                    self.connections.add(connection)
+
+                    logger.info("Starting interaction %r %r", *addr)
+                    try:
+                        await connection.run_application()
+                    finally:
+                        self.connections.remove(connection)
+                        logger.info("Stopping interaction %r %r", *addr)
+            except BaseException as e:
+                print("Got %s" % type(e).__name__, e)
+                import traceback
+
+                traceback.print_exc()
             finally:
-                self.connections.remove(connection)
                 self._application_tasks.remove(task)
-                logger.info("Stopping interaction %r %r", *addr)
 
         task = get_event_loop().create_task(run())
         self._application_tasks.append(task)

prompt_toolkit/input/__init__.py

@@ -1,9 +1,10 @@
-from .base import DummyInput, Input
+from .base import DummyInput, Input, PipeInput
 from .defaults import create_input, create_pipe_input
 
 __all__ = [
     # Base.
     "Input",
+    "PipeInput",
     "DummyInput",
     # Defaults.
     "create_input",

prompt_toolkit/input/base.py

@@ -9,6 +9,7 @@
 
 __all__ = [
     "Input",
+    "PipeInput",
     "DummyInput",
 ]
 

prompt_toolkit/input/defaults.py

@@ -1,5 +1,5 @@
 import sys
-from typing import Optional, TextIO
+from typing import ContextManager, Optional, TextIO
 
 from prompt_toolkit.utils import is_windows
 
@@ -48,16 +48,24 @@ def create_input(
         return Vt100Input(stdin)
 
 
-def create_pipe_input() -> PipeInput:
+def create_pipe_input() -> ContextManager[PipeInput]:
     """
     Create an input pipe.
     This is mostly useful for unit testing.
+
+    Usage::
+
+        with create_pipe_input() as input:
+            input.send_text('inputdata')
+
+    Breaking change: In prompt_toolkit 3.0.28 and earlier, this was returning
+    the `PipeInput` directly, rather than through a context manager.
     """
     if is_windows():
         from .win32_pipe import Win32PipeInput
 
-        return Win32PipeInput()
+        return Win32PipeInput.create()
     else:
         from .posix_pipe import PosixPipeInput
 
-        return PosixPipeInput()
+        return PosixPipeInput.create()

prompt_toolkit/input/posix_pipe.py

@@ -1,5 +1,6 @@
 import os
-from typing import ContextManager, TextIO, cast
+from contextlib import contextmanager
+from typing import ContextManager, Iterator, TextIO, cast
 
 from ..utils import DummyContext
 from .base import PipeInput
@@ -10,6 +11,36 @@
 ]
 
 
+class _Pipe:
+    "Wrapper around os.pipe, that ensures we don't double close any end."
+
+    def __init__(self) -> None:
+        self.read_fd, self.write_fd = os.pipe()
+        self._read_closed = False
+        self._write_closed = False
+
+    def close_read(self) -> None:
+        "Close read-end if not yet closed."
+        if self._read_closed:
+            return
+
+        os.close(self.read_fd)
+        self._read_closed = True
+
+    def close_write(self) -> None:
+        "Close write-end if not yet closed."
+        if self._write_closed:
+            return
+
+        os.close(self.write_fd)
+        self._write_closed = True
+
+    def close(self) -> None:
+        "Close both read and write ends."
+        self.close_read()
+        self.close_write()
+
+
 class PosixPipeInput(Vt100Input, PipeInput):
     """
     Input that is send through a pipe.
@@ -18,14 +49,15 @@ class PosixPipeInput(Vt100Input, PipeInput):
 
     Usage::
 
-        input = PosixPipeInput()
-        input.send_text('inputdata')
+        with PosixPipeInput.create() as input:
+            input.send_text('inputdata')
     """
 
     _id = 0
 
-    def __init__(self, text: str = "") -> None:
-        self._r, self._w = os.pipe()
+    def __init__(self, _pipe: _Pipe, _text: str = "") -> None:
+        # Private constructor. Users should use the public `.create()` method.
+        self.pipe = _pipe
 
         class Stdin:
             encoding = "utf-8"
@@ -34,21 +66,30 @@ def isatty(stdin) -> bool:
                 return True
 
             def fileno(stdin) -> int:
-                return self._r
+                return self.pipe.read_fd
 
         super().__init__(cast(TextIO, Stdin()))
-        self.send_text(text)
+        self.send_text(_text)
 
         # Identifier for every PipeInput for the hash.
         self.__class__._id += 1
         self._id = self.__class__._id
 
+    @classmethod
+    @contextmanager
+    def create(cls, text: str = "") -> Iterator["PosixPipeInput"]:
+        pipe = _Pipe()
+        try:
+            yield PosixPipeInput(_pipe=pipe, _text=text)
+        finally:
+            pipe.close()
+
     def send_bytes(self, data: bytes) -> None:
-        os.write(self._w, data)
+        os.write(self.pipe.write_fd, data)
 
     def send_text(self, data: str) -> None:
         "Send text to the input."
-        os.write(self._w, data.encode("utf-8"))
+        os.write(self.pipe.write_fd, data.encode("utf-8"))
 
     def raw_mode(self) -> ContextManager[None]:
         return DummyContext()
@@ -58,12 +99,11 @@ def cooked_mode(self) -> ContextManager[None]:
 
     def close(self) -> None:
         "Close pipe fds."
-        os.close(self._r)
-        os.close(self._w)
-
-        # We should assign `None` to 'self._r` and 'self._w',
-        # The event loop still needs to know the the fileno for this input in order
-        # to properly remove it from the selectors.
+        # Only close the write-end of the pipe. This will unblock the reader
+        # callback (in vt100.py > _attached_input), which eventually will raise
+        # `EOFError`. If we'd also close the read-end, then the event loop
+        # won't wake up the corresponding callback because of this.
+        self.pipe.close_write()
 
     def typeahead_hash(self) -> str:
         """