Improved handling of stdout objects that don't have a 'buffer' attribute.

For instance, when using `renderer_print_formatted_text` in a Jupyter Notebook.

Author: jonathanslenders

src/prompt_toolkit/contrib/ssh/server.py

@@ -85,9 +85,8 @@ async def _interact(self) -> None:
 
         term = self._chan.get_terminal_type()
 
-        self._output = Vt100_Output(
-            self.stdout, self._get_size, term=term, write_binary=False
-        )
+        self._output = Vt100_Output(self.stdout, self._get_size, term=term)
+
         with create_pipe_input() as self._input:
             with create_app_session(input=self._input, output=self._output) as session:
                 self.app_session = session

src/prompt_toolkit/contrib/telnet/server.py

@@ -169,9 +169,7 @@ def size_received(rows: int, columns: int) -> None:
 
         def ttype_received(ttype: str) -> None:
             """TelnetProtocolParser 'ttype_received' callback"""
-            self.vt100_output = Vt100_Output(
-                self.stdout, get_size, term=ttype, write_binary=False
-            )
+            self.vt100_output = Vt100_Output(self.stdout, get_size, term=ttype)
             self._ready.set()
 
         self.parser = TelnetProtocolParser(data_received, size_received, ttype_received)

src/prompt_toolkit/output/flush_stdout.py

@@ -7,7 +7,15 @@
 __all__ = ["flush_stdout"]
 
 
-def flush_stdout(stdout: TextIO, data: str, write_binary: bool) -> None:
+def flush_stdout(stdout: TextIO, data: str) -> None:
+    # If the IO object has an `encoding` and `buffer` attribute, it means that
+    # we can access the underlying BinaryIO object and write into it in binary
+    # mode. This is preferred if possible.
+    # NOTE: When used in a Jupyter notebook, don't write binary.
+    #       `ipykernel.iostream.OutStream` has an `encoding` attribute, but not
+    #       a `buffer` attribute, so we can't write binary in it.
+    has_binary_io = hasattr(stdout, "encoding") and hasattr(stdout, "buffer")
+
     try:
         # Ensure that `stdout` is made blocking when writing into it.
         # Otherwise, when uvloop is activated (which makes stdout
@@ -20,14 +28,8 @@ def flush_stdout(stdout: TextIO, data: str, write_binary: bool) -> None:
             # My Arch Linux installation of july 2015 reported 'ANSI_X3.4-1968'
             # for sys.stdout.encoding in xterm.
             out: IO[bytes]
-            if write_binary:
-                if hasattr(stdout, "buffer"):
-                    out = stdout.buffer
-                else:
-                    # IO[bytes] was given to begin with.
-                    # (Used in the unit tests, for instance.)
-                    out = cast(IO[bytes], stdout)
-                out.write(data.encode(stdout.encoding or "utf-8", "replace"))
+            if has_binary_io:
+                stdout.buffer.write(data.encode(stdout.encoding or "utf-8", "replace"))
             else:
                 stdout.write(data)
 

src/prompt_toolkit/output/plain_text.py

@@ -23,14 +23,10 @@ class PlainTextOutput(Output):
     formatting.)
     """
 
-    def __init__(self, stdout: TextIO, write_binary: bool = True) -> None:
+    def __init__(self, stdout: TextIO) -> None:
         assert all(hasattr(stdout, a) for a in ("write", "flush"))
 
-        if write_binary:
-            assert hasattr(stdout, "encoding")
-
         self.stdout: TextIO = stdout
-        self.write_binary = write_binary
         self._buffer: List[str] = []
 
     def fileno(self) -> int:
@@ -58,7 +54,7 @@ def flush(self) -> None:
 
         data = "".join(self._buffer)
         self._buffer = []
-        flush_stdout(self.stdout, data, write_binary=self.write_binary)
+        flush_stdout(self.stdout, data)
 
     def erase_screen(self) -> None:
         pass

src/prompt_toolkit/output/vt100.py

@@ -400,8 +400,6 @@ class Vt100_Output(Output):
     :param get_size: A callable which returns the `Size` of the output terminal.
     :param stdout: Any object with has a `write` and `flush` method + an 'encoding' property.
     :param term: The terminal environment variable. (xterm, xterm-256color, linux, ...)
-    :param write_binary: Encode the output before writing it. If `True` (the
-        default), the `stdout` object is supposed to expose an `encoding` attribute.
     """
 
     # For the error messages. Only display "Output is not a terminal" once per
@@ -413,19 +411,14 @@ def __init__(
         stdout: TextIO,
         get_size: Callable[[], Size],
         term: Optional[str] = None,
-        write_binary: bool = True,
         default_color_depth: Optional[ColorDepth] = None,
         enable_bell: bool = True,
     ) -> None:
 
         assert all(hasattr(stdout, a) for a in ("write", "flush"))
 
-        if write_binary:
-            assert hasattr(stdout, "encoding")
-
         self._buffer: List[str] = []
         self.stdout: TextIO = stdout
-        self.write_binary = write_binary
         self.default_color_depth = default_color_depth
         self._get_size = get_size
         self.term = term
@@ -699,7 +692,7 @@ def flush(self) -> None:
         data = "".join(self._buffer)
         self._buffer = []
 
-        flush_stdout(self.stdout, data, write_binary=self.write_binary)
+        flush_stdout(self.stdout, data)
 
     def ask_for_cpr(self) -> None:
         """

tests/test_print_formatted_text.py

@@ -12,7 +12,6 @@
 
 class _Capture:
     "Emulate an stdout object."
-    encoding = "utf-8"
 
     def __init__(self):
         self._data = []
@@ -22,7 +21,7 @@ def write(self, data):
 
     @property
     def data(self):
-        return b"".join(self._data)
+        return "".join(self._data)
 
     def flush(self):
         pass
@@ -40,15 +39,15 @@ def fileno(self):
 def test_print_formatted_text():
     f = _Capture()
     pt_print([("", "hello"), ("", "world")], file=f)
-    assert b"hello" in f.data
-    assert b"world" in f.data
+    assert "hello" in f.data
+    assert "world" in f.data
 
 
 @pytest.mark.skipif(is_windows(), reason="Doesn't run on Windows yet.")
 def test_print_formatted_text_backslash_r():
     f = _Capture()
     pt_print("hello\r\n", file=f)
-    assert b"hello" in f.data
+    assert "hello" in f.data
 
 
 @pytest.mark.skipif(is_windows(), reason="Doesn't run on Windows yet.")
@@ -70,8 +69,8 @@ def test_formatted_text_with_style():
     # NOTE: We pass the default (8bit) color depth, so that the unit tests
     #       don't start failing when environment variables change.
     pt_print(tokens, style=style, file=f, color_depth=ColorDepth.DEFAULT)
-    assert b"\x1b[0;38;5;197mHello" in f.data
-    assert b"\x1b[0;38;5;83;3mworld" in f.data
+    assert "\x1b[0;38;5;197mHello" in f.data
+    assert "\x1b[0;38;5;83;3mworld" in f.data
 
 
 @pytest.mark.skipif(is_windows(), reason="Doesn't run on Windows yet.")
@@ -87,5 +86,5 @@ def test_html_with_style():
 
     assert (
         f.data
-        == b"\x1b[0m\x1b[?7h\x1b[0;32mhello\x1b[0m \x1b[0;1mworld\x1b[0m\r\n\x1b[0m"
+        == "\x1b[0m\x1b[?7h\x1b[0;32mhello\x1b[0m \x1b[0;1mworld\x1b[0m\r\n\x1b[0m"
     )