io.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. # encoding: utf-8
  2. """
  3. IO related utilities.
  4. """
  5. # Copyright (c) IPython Development Team.
  6. # Distributed under the terms of the Modified BSD License.
  7. import atexit
  8. import os
  9. import sys
  10. import tempfile
  11. from pathlib import Path
  12. from warnings import warn
  13. from IPython.utils.decorators import undoc
  14. from .capture import CapturedIO, capture_output
  15. from io import StringIO
  16. from typing import Union
  17. class Tee:
  18. """A class to duplicate an output stream to stdout/err.
  19. This works in a manner very similar to the Unix 'tee' command.
  20. When the object is closed or deleted, it closes the original file given to
  21. it for duplication.
  22. """
  23. # Inspired by:
  24. # http://mail.python.org/pipermail/python-list/2007-May/442737.html
  25. def __init__(self, file_or_name: Union[str, StringIO], mode: str="w", channel: str='stdout'):
  26. """Construct a new Tee object.
  27. Parameters
  28. ----------
  29. file_or_name : filename or open filehandle (writable)
  30. File that will be duplicated
  31. mode : optional, valid mode for open().
  32. If a filename was give, open with this mode.
  33. channel : str, one of ['stdout', 'stderr']
  34. """
  35. if channel not in ['stdout', 'stderr']:
  36. raise ValueError('Invalid channel spec %s' % channel)
  37. if hasattr(file_or_name, 'write') and hasattr(file_or_name, 'seek'):
  38. self.file = file_or_name
  39. else:
  40. encoding = None if "b" in mode else "utf-8"
  41. self.file = open(file_or_name, mode, encoding=encoding)
  42. self.channel = channel
  43. self.ostream = getattr(sys, channel)
  44. setattr(sys, channel, self)
  45. self._closed = False
  46. def close(self):
  47. """Close the file and restore the channel."""
  48. self.flush()
  49. setattr(sys, self.channel, self.ostream)
  50. self.file.close()
  51. self._closed = True
  52. def write(self, data):
  53. """Write data to both channels."""
  54. self.file.write(data)
  55. self.ostream.write(data)
  56. self.ostream.flush()
  57. def flush(self):
  58. """Flush both channels."""
  59. self.file.flush()
  60. self.ostream.flush()
  61. def __del__(self):
  62. if not self._closed:
  63. self.close()
  64. def isatty(self):
  65. return False
  66. def ask_yes_no(prompt, default=None, interrupt=None):
  67. """Asks a question and returns a boolean (y/n) answer.
  68. If default is given (one of 'y','n'), it is used if the user input is
  69. empty. If interrupt is given (one of 'y','n'), it is used if the user
  70. presses Ctrl-C. Otherwise the question is repeated until an answer is
  71. given.
  72. An EOF is treated as the default answer. If there is no default, an
  73. exception is raised to prevent infinite loops.
  74. Valid answers are: y/yes/n/no (match is not case sensitive)."""
  75. answers = {'y':True,'n':False,'yes':True,'no':False}
  76. ans = None
  77. while ans not in answers.keys():
  78. try:
  79. ans = input(prompt+' ').lower()
  80. if not ans: # response was an empty string
  81. ans = default
  82. except KeyboardInterrupt:
  83. if interrupt:
  84. ans = interrupt
  85. print("\r")
  86. except EOFError:
  87. if default in answers.keys():
  88. ans = default
  89. print()
  90. else:
  91. raise
  92. return answers[ans]
  93. def temp_pyfile(src: str, ext: str='.py') -> str:
  94. """Make a temporary python file, return filename and filehandle.
  95. Parameters
  96. ----------
  97. src : string or list of strings (no need for ending newlines if list)
  98. Source code to be written to the file.
  99. ext : optional, string
  100. Extension for the generated file.
  101. Returns
  102. -------
  103. (filename, open filehandle)
  104. It is the caller's responsibility to close the open file and unlink it.
  105. """
  106. fname = tempfile.mkstemp(ext)[1]
  107. with open(Path(fname), "w", encoding="utf-8") as f:
  108. f.write(src)
  109. f.flush()
  110. return fname