/jishaku/shell.py

https://github.com/Gorialis/jishaku · Python · 147 lines · 116 code · 8 blank · 23 comment · 2 complexity · 379d008356b64252c6c90be721afbc72 MD5 · raw file

  1. # -*- coding: utf-8 -*-
  2. """
  3. jishaku.shell
  4. ~~~~~~~~~~~~~
  5. Tools related to interacting directly with the shell.
  6. :copyright: (c) 2021 Devon (Gorialis) R
  7. :license: MIT, see LICENSE for more details.
  8. """
  9. import asyncio
  10. import os
  11. import pathlib
  12. import re
  13. import subprocess
  14. import sys
  15. import time
  16. SHELL = os.getenv("SHELL") or "/bin/bash"
  17. WINDOWS = sys.platform == "win32"
  18. def background_reader(stream, loop: asyncio.AbstractEventLoop, callback):
  19. """
  20. Reads a stream and forwards each line to an async callback.
  21. """
  22. for line in iter(stream.readline, b''):
  23. loop.call_soon_threadsafe(loop.create_task, callback(line))
  24. class ShellReader:
  25. """
  26. A class that passively reads from a shell and buffers results for read.
  27. Example
  28. -------
  29. .. code:: python3
  30. # reader should be in a with statement to ensure it is properly closed
  31. with ShellReader('echo one; sleep 5; echo two') as reader:
  32. # prints 'one', then 'two' after 5 seconds
  33. async for x in reader:
  34. print(x)
  35. """
  36. def __init__(self, code: str, timeout: int = 120, loop: asyncio.AbstractEventLoop = None):
  37. if WINDOWS:
  38. # Check for powershell
  39. if pathlib.Path(r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe").exists():
  40. sequence = ['powershell', code]
  41. self.ps1 = "PS >"
  42. self.highlight = "powershell"
  43. else:
  44. sequence = ['cmd', '/c', code]
  45. self.ps1 = "cmd >"
  46. self.highlight = "cmd"
  47. else:
  48. sequence = [SHELL, '-c', code]
  49. self.ps1 = "$"
  50. self.highlight = "sh"
  51. self.process = subprocess.Popen(sequence, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # pylint: disable=consider-using-with
  52. self.close_code = None
  53. self.loop = loop or asyncio.get_event_loop()
  54. self.timeout = timeout
  55. self.stdout_task = self.make_reader_task(self.process.stdout, self.stdout_handler)
  56. self.stderr_task = self.make_reader_task(self.process.stderr, self.stderr_handler)
  57. self.queue = asyncio.Queue(maxsize=250)
  58. @property
  59. def closed(self):
  60. """
  61. Are both tasks done, indicating there is no more to read?
  62. """
  63. return self.stdout_task.done() and self.stderr_task.done()
  64. async def executor_wrapper(self, *args, **kwargs):
  65. """
  66. Call wrapper for stream reader.
  67. """
  68. return await self.loop.run_in_executor(None, *args, **kwargs)
  69. def make_reader_task(self, stream, callback):
  70. """
  71. Create a reader executor task for a stream.
  72. """
  73. return self.loop.create_task(self.executor_wrapper(background_reader, stream, self.loop, callback))
  74. @staticmethod
  75. def clean_bytes(line):
  76. """
  77. Cleans a byte sequence of shell directives and decodes it.
  78. """
  79. text = line.decode('utf-8').replace('\r', '').strip('\n')
  80. return re.sub(r'\x1b[^m]*m', '', text).replace("``", "`\u200b`").strip('\n')
  81. async def stdout_handler(self, line):
  82. """
  83. Handler for this class for stdout.
  84. """
  85. await self.queue.put(self.clean_bytes(line))
  86. async def stderr_handler(self, line):
  87. """
  88. Handler for this class for stderr.
  89. """
  90. await self.queue.put(self.clean_bytes(b'[stderr] ' + line))
  91. def __enter__(self):
  92. return self
  93. def __exit__(self, *args):
  94. self.process.kill()
  95. self.process.terminate()
  96. self.close_code = self.process.wait(timeout=0.5)
  97. def __aiter__(self):
  98. return self
  99. async def __anext__(self):
  100. last_output = time.perf_counter()
  101. while not self.closed or not self.queue.empty():
  102. try:
  103. item = await asyncio.wait_for(self.queue.get(), timeout=1)
  104. except asyncio.TimeoutError as exception:
  105. if time.perf_counter() - last_output >= self.timeout:
  106. raise exception
  107. else:
  108. last_output = time.perf_counter()
  109. return item
  110. raise StopAsyncIteration()