flush_stdout.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. from __future__ import annotations
  2. import errno
  3. import os
  4. import sys
  5. from contextlib import contextmanager
  6. from typing import IO, Iterator, TextIO
  7. __all__ = ["flush_stdout"]
  8. def flush_stdout(stdout: TextIO, data: str) -> None:
  9. # If the IO object has an `encoding` and `buffer` attribute, it means that
  10. # we can access the underlying BinaryIO object and write into it in binary
  11. # mode. This is preferred if possible.
  12. # NOTE: When used in a Jupyter notebook, don't write binary.
  13. # `ipykernel.iostream.OutStream` has an `encoding` attribute, but not
  14. # a `buffer` attribute, so we can't write binary in it.
  15. has_binary_io = hasattr(stdout, "encoding") and hasattr(stdout, "buffer")
  16. try:
  17. # Ensure that `stdout` is made blocking when writing into it.
  18. # Otherwise, when uvloop is activated (which makes stdout
  19. # non-blocking), and we write big amounts of text, then we get a
  20. # `BlockingIOError` here.
  21. with _blocking_io(stdout):
  22. # (We try to encode ourself, because that way we can replace
  23. # characters that don't exist in the character set, avoiding
  24. # UnicodeEncodeError crashes. E.g. u'\xb7' does not appear in 'ascii'.)
  25. # My Arch Linux installation of july 2015 reported 'ANSI_X3.4-1968'
  26. # for sys.stdout.encoding in xterm.
  27. if has_binary_io:
  28. stdout.buffer.write(data.encode(stdout.encoding or "utf-8", "replace"))
  29. else:
  30. stdout.write(data)
  31. stdout.flush()
  32. except OSError as e:
  33. if e.args and e.args[0] == errno.EINTR:
  34. # Interrupted system call. Can happen in case of a window
  35. # resize signal. (Just ignore. The resize handler will render
  36. # again anyway.)
  37. pass
  38. elif e.args and e.args[0] == 0:
  39. # This can happen when there is a lot of output and the user
  40. # sends a KeyboardInterrupt by pressing Control-C. E.g. in
  41. # a Python REPL when we execute "while True: print('test')".
  42. # (The `ptpython` REPL uses this `Output` class instead of
  43. # `stdout` directly -- in order to be network transparent.)
  44. # So, just ignore.
  45. pass
  46. else:
  47. raise
  48. @contextmanager
  49. def _blocking_io(io: IO[str]) -> Iterator[None]:
  50. """
  51. Ensure that the FD for `io` is set to blocking in here.
  52. """
  53. if sys.platform == "win32":
  54. # On Windows, the `os` module doesn't have a `get/set_blocking`
  55. # function.
  56. yield
  57. return
  58. try:
  59. fd = io.fileno()
  60. blocking = os.get_blocking(fd)
  61. except: # noqa
  62. # Failed somewhere.
  63. # `get_blocking` can raise `OSError`.
  64. # The io object can raise `AttributeError` when no `fileno()` method is
  65. # present if we're not a real file object.
  66. blocking = True # Assume we're good, and don't do anything.
  67. try:
  68. # Make blocking if we weren't blocking yet.
  69. if not blocking:
  70. os.set_blocking(fd, True)
  71. yield
  72. finally:
  73. # Restore original blocking mode.
  74. if not blocking:
  75. os.set_blocking(fd, blocking)