PageRenderTime 51ms CodeModel.GetById 21ms RepoModel.GetById 0ms app.codeStats 0ms

/xonsh/teepty.py

https://gitlab.com/vectorci/gitsome
Python | 331 lines | 293 code | 12 blank | 26 comment | 11 complexity | 9767d4acfb3a9365ead66d4e0b1bdcee MD5 | raw file
  1. """This implements a psuedo-TTY that tees its output into a Python buffer.
  2. This file was forked from a version distibuted under an MIT license and
  3. Copyright (c) 2011 Joshua D. Bartlett.
  4. See http://sqizit.bartletts.id.au/2011/02/14/pseudo-terminals-in-python/ for
  5. more information.
  6. """
  7. import io
  8. import re
  9. import os
  10. import sys
  11. import tty
  12. import pty
  13. import time
  14. import array
  15. import fcntl
  16. import select
  17. import signal
  18. import termios
  19. import tempfile
  20. import threading
  21. # The following escape codes are xterm codes.
  22. # See http://rtfm.etla.org/xterm/ctlseq.html for more.
  23. MODE_NUMS = ('1049', '47', '1047')
  24. START_ALTERNATE_MODE = frozenset('\x1b[?{0}h'.format(i).encode() for i in MODE_NUMS)
  25. END_ALTERNATE_MODE = frozenset('\x1b[?{0}l'.format(i).encode() for i in MODE_NUMS)
  26. ALTERNATE_MODE_FLAGS = tuple(START_ALTERNATE_MODE) + tuple(END_ALTERNATE_MODE)
  27. RE_HIDDEN = re.compile(b'(\001.*?\002)')
  28. RE_COLOR = re.compile(b'\033\[\d+;?\d*m')
  29. def _findfirst(s, substrs):
  30. """Finds whichever of the given substrings occurs first in the given string
  31. and returns that substring, or returns None if no such strings occur.
  32. """
  33. i = len(s)
  34. result = None
  35. for substr in substrs:
  36. pos = s.find(substr)
  37. if -1 < pos < i:
  38. i = pos
  39. result = substr
  40. return i, result
  41. def _on_main_thread():
  42. """Checks if we are on the main thread or not. Duplicated from xonsh.tools
  43. here so that this module only relies on the Python standrd library.
  44. """
  45. return threading.current_thread() is threading.main_thread()
  46. def _find_error_code(e):
  47. """Gets the approriate error code for an exception e, see
  48. http://tldp.org/LDP/abs/html/exitcodes.html for exit codes.
  49. """
  50. if isinstance(e, PermissionError):
  51. code = 126
  52. elif isinstance(e, FileNotFoundError):
  53. code = 127
  54. else:
  55. code = 1
  56. return code
  57. class TeePTY(object):
  58. """This class is a pseudo terminal that tees the stdout and stderr into a buffer."""
  59. def __init__(self, bufsize=1024, remove_color=True, encoding='utf-8',
  60. errors='strict'):
  61. """
  62. Parameters
  63. ----------
  64. bufsize : int, optional
  65. The buffer size to read from the root terminal to/from the tee'd terminal.
  66. remove_color : bool, optional
  67. Removes color codes from the tee'd buffer, though not the TTY.
  68. encoding : str, optional
  69. The encoding to use when decoding into a str.
  70. errors : str, optional
  71. The encoding error flag to use when decoding into a str.
  72. """
  73. self.bufsize = bufsize
  74. self.pid = self.master_fd = None
  75. self._in_alt_mode = False
  76. self.remove_color = remove_color
  77. self.encoding = encoding
  78. self.errors = errors
  79. self.buffer = io.BytesIO()
  80. self.returncode = None
  81. self._temp_stdin = None
  82. def __str__(self):
  83. return self.buffer.getvalue().decode(encoding=self.encoding,
  84. errors=self.errors)
  85. def __del__(self):
  86. if self._temp_stdin is not None:
  87. self._temp_stdin.close()
  88. self._temp_stdin = None
  89. def spawn(self, argv=None, env=None, stdin=None, delay=None):
  90. """Create a spawned process. Based on the code for pty.spawn().
  91. This cannot be used except from the main thread.
  92. Parameters
  93. ----------
  94. argv : list of str, optional
  95. Arguments to pass in as subprocess. In None, will execute $SHELL.
  96. env : Mapping, optional
  97. Environment to pass execute in.
  98. delay : float, optional
  99. Delay timing before executing process if piping in data. The value
  100. is passed into time.sleep() so it is in [seconds]. If delay is None,
  101. its value will attempted to be looked up from the environment
  102. variable $TEEPTY_PIPE_DELAY, from the passed in env or os.environ.
  103. If not present or not positive valued, no delay is used.
  104. Returns
  105. -------
  106. returncode : int
  107. Return code for the spawned process.
  108. """
  109. assert self.master_fd is None
  110. self._in_alt_mode = False
  111. if not argv:
  112. argv = [os.environ.get('SHELL', 'sh')]
  113. argv = self._put_stdin_in_argv(argv, stdin)
  114. pid, master_fd = pty.fork()
  115. self.pid = pid
  116. self.master_fd = master_fd
  117. if pid == pty.CHILD:
  118. # determine if a piping delay is needed.
  119. if self._temp_stdin is not None:
  120. self._delay_for_pipe(env=env, delay=delay)
  121. # ok, go
  122. try:
  123. if env is None:
  124. os.execvp(argv[0], argv)
  125. else:
  126. os.execvpe(argv[0], argv, env)
  127. except OSError as e:
  128. os._exit(_find_error_code(e))
  129. else:
  130. self._pipe_stdin(stdin)
  131. on_main_thread = _on_main_thread()
  132. if on_main_thread:
  133. old_handler = signal.signal(signal.SIGWINCH, self._signal_winch)
  134. try:
  135. mode = tty.tcgetattr(pty.STDIN_FILENO)
  136. tty.setraw(pty.STDIN_FILENO)
  137. restore = True
  138. except tty.error: # This is the same as termios.error
  139. restore = False
  140. self._init_fd()
  141. try:
  142. self._copy()
  143. except (IOError, OSError):
  144. if restore:
  145. tty.tcsetattr(pty.STDIN_FILENO, tty.TCSAFLUSH, mode)
  146. _, self.returncode = os.waitpid(pid, 0)
  147. os.close(master_fd)
  148. self.master_fd = None
  149. self._in_alt_mode = False
  150. if on_main_thread:
  151. signal.signal(signal.SIGWINCH, old_handler)
  152. return self.returncode
  153. def _init_fd(self):
  154. """Called once when the pty is first set up."""
  155. self._set_pty_size()
  156. def _signal_winch(self, signum, frame):
  157. """Signal handler for SIGWINCH - window size has changed."""
  158. self._set_pty_size()
  159. def _set_pty_size(self):
  160. """Sets the window size of the child pty based on the window size of
  161. our own controlling terminal.
  162. """
  163. assert self.master_fd is not None
  164. # Get the terminal size of the real terminal, set it on the
  165. # pseudoterminal.
  166. buf = array.array('h', [0, 0, 0, 0])
  167. fcntl.ioctl(pty.STDOUT_FILENO, termios.TIOCGWINSZ, buf, True)
  168. fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, buf)
  169. def _copy(self):
  170. """Main select loop. Passes all data to self.master_read() or self.stdin_read().
  171. """
  172. assert self.master_fd is not None
  173. master_fd = self.master_fd
  174. bufsize = self.bufsize
  175. while True:
  176. try:
  177. rfds, wfds, xfds = select.select([master_fd, pty.STDIN_FILENO], [], [])
  178. except OSError as e:
  179. if e.errno == 4: # Interrupted system call.
  180. continue # This happens at terminal resize.
  181. if master_fd in rfds:
  182. data = os.read(master_fd, bufsize)
  183. self.write_stdout(data)
  184. if pty.STDIN_FILENO in rfds:
  185. data = os.read(pty.STDIN_FILENO, bufsize)
  186. self.write_stdin(data)
  187. def _sanatize_data(self, data):
  188. i, flag = _findfirst(data, ALTERNATE_MODE_FLAGS)
  189. if flag is None and self._in_alt_mode:
  190. return b''
  191. elif flag is not None:
  192. if flag in START_ALTERNATE_MODE:
  193. # This code is executed when the child process switches the terminal into
  194. # alternate mode. The line below assumes that the user has opened vim,
  195. # less, or similar, and writes writes to stdin.
  196. d0 = data[:i]
  197. self._in_alt_mode = True
  198. d1 = self._sanatize_data(data[i+len(flag):])
  199. data = d0 + d1
  200. elif flag in END_ALTERNATE_MODE:
  201. # This code is executed when the child process switches the terminal back
  202. # out of alternate mode. The line below assumes that the user has
  203. # returned to the command prompt.
  204. self._in_alt_mode = False
  205. data = self._sanatize_data(data[i+len(flag):])
  206. data = RE_HIDDEN.sub(b'', data)
  207. if self.remove_color:
  208. data = RE_COLOR.sub(b'', data)
  209. return data
  210. def write_stdout(self, data):
  211. """Writes to stdout as if the child process had written the data (bytes)."""
  212. os.write(pty.STDOUT_FILENO, data) # write to real terminal
  213. # tee to buffer
  214. data = self._sanatize_data(data)
  215. if len(data) > 0:
  216. self.buffer.write(data)
  217. def write_stdin(self, data):
  218. """Writes to the child process from its controlling terminal."""
  219. master_fd = self.master_fd
  220. assert master_fd is not None
  221. while len(data) > 0:
  222. n = os.write(master_fd, data)
  223. data = data[n:]
  224. def _stdin_filename(self, stdin):
  225. if stdin is None:
  226. rtn = None
  227. elif isinstance(stdin, io.FileIO) and os.path.isfile(stdin.name):
  228. rtn = stdin.name
  229. elif isinstance(stdin, (io.BufferedIOBase, str, bytes)):
  230. self._temp_stdin = tsi = tempfile.NamedTemporaryFile()
  231. rtn = tsi.name
  232. else:
  233. raise ValueError('stdin not understood {0!r}'.format(stdin))
  234. return rtn
  235. def _put_stdin_in_argv(self, argv, stdin):
  236. stdin_filename = self._stdin_filename(stdin)
  237. if stdin_filename is None:
  238. return argv
  239. argv = list(argv)
  240. # a lone dash '-' argument means stdin
  241. if argv.count('-') == 0:
  242. argv.append(stdin_filename)
  243. else:
  244. argv[argv.index('-')] = stdin_filename
  245. return argv
  246. def _pipe_stdin(self, stdin):
  247. if stdin is None or isinstance(stdin, io.FileIO):
  248. return None
  249. tsi = self._temp_stdin
  250. bufsize = self.bufsize
  251. if isinstance(stdin, io.BufferedIOBase):
  252. buf = stdin.read(bufsize)
  253. while len(buf) != 0:
  254. tsi.write(buf)
  255. tsi.flush()
  256. buf = stdin.read(bufsize)
  257. elif isinstance(stdin, (str, bytes)):
  258. raw = stdin.encode() if isinstance(stdin, str) else stdin
  259. for i in range((len(raw)//bufsize) + 1):
  260. tsi.write(raw[i*bufsize:(i + 1)*bufsize])
  261. tsi.flush()
  262. else:
  263. raise ValueError('stdin not understood {0!r}'.format(stdin))
  264. def _delay_for_pipe(self, env=None, delay=None):
  265. # This delay is sometimes needed because the temporary stdin file that
  266. # is being written (the pipe) may not have even hits its first flush()
  267. # call by the time the spawned process starts up and determines there
  268. # is nothing in the file. The spawn can thus exit, without doing any
  269. # real work. Consider the case of piping something into grep:
  270. #
  271. # $ ps aux | grep root
  272. #
  273. # grep will exit on EOF and so there is a race between the buffersize
  274. # and flushing the temporary file and grep. However, this race is not
  275. # always meaningful. Pagers, for example, update when the file is written
  276. # to. So what is important is that we start the spawned process ASAP:
  277. #
  278. # $ ps aux | less
  279. #
  280. # So there is a push-and-pull between the the competing objectives of
  281. # not blocking and letting the spawned process have enough to work with
  282. # such that it doesn't exit prematurely. Unfortunately, there is no
  283. # way to know a priori how big the file is, how long the spawned process
  284. # will run for, etc. Thus as user-definable delay let's the user
  285. # find something that works for them.
  286. if delay is None:
  287. delay = (env or os.environ).get('TEEPTY_PIPE_DELAY', -1.0)
  288. delay = float(delay)
  289. if 0.0 < delay:
  290. time.sleep(delay)
  291. if __name__ == '__main__':
  292. tpty = TeePTY()
  293. tpty.spawn(sys.argv[1:])
  294. print('-=-'*10)
  295. print(tpty.buffer.getvalue())
  296. print('-=-'*10)
  297. print(tpty)
  298. print('-=-'*10)
  299. print('Returned with status {0}'.format(tpty.returncode))