__main__.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128
  1. import ast
  2. import asyncio
  3. import code
  4. import concurrent.futures
  5. import contextvars
  6. import inspect
  7. import sys
  8. import threading
  9. import types
  10. import warnings
  11. from . import futures
  12. class AsyncIOInteractiveConsole(code.InteractiveConsole):
  13. def __init__(self, locals, loop):
  14. super().__init__(locals)
  15. self.compile.compiler.flags |= ast.PyCF_ALLOW_TOP_LEVEL_AWAIT
  16. self.loop = loop
  17. self.context = contextvars.copy_context()
  18. def runcode(self, code):
  19. future = concurrent.futures.Future()
  20. def callback():
  21. global repl_future
  22. global repl_future_interrupted
  23. repl_future = None
  24. repl_future_interrupted = False
  25. func = types.FunctionType(code, self.locals)
  26. try:
  27. coro = func()
  28. except SystemExit:
  29. raise
  30. except KeyboardInterrupt as ex:
  31. repl_future_interrupted = True
  32. future.set_exception(ex)
  33. return
  34. except BaseException as ex:
  35. future.set_exception(ex)
  36. return
  37. if not inspect.iscoroutine(coro):
  38. future.set_result(coro)
  39. return
  40. try:
  41. repl_future = self.loop.create_task(coro, context=self.context)
  42. futures._chain_future(repl_future, future)
  43. except BaseException as exc:
  44. future.set_exception(exc)
  45. loop.call_soon_threadsafe(callback, context=self.context)
  46. try:
  47. return future.result()
  48. except SystemExit:
  49. raise
  50. except BaseException:
  51. if repl_future_interrupted:
  52. self.write("\nKeyboardInterrupt\n")
  53. else:
  54. self.showtraceback()
  55. class REPLThread(threading.Thread):
  56. def run(self):
  57. try:
  58. banner = (
  59. f'asyncio REPL {sys.version} on {sys.platform}\n'
  60. f'Use "await" directly instead of "asyncio.run()".\n'
  61. f'Type "help", "copyright", "credits" or "license" '
  62. f'for more information.\n'
  63. f'{getattr(sys, "ps1", ">>> ")}import asyncio'
  64. )
  65. console.interact(
  66. banner=banner,
  67. exitmsg='exiting asyncio REPL...')
  68. finally:
  69. warnings.filterwarnings(
  70. 'ignore',
  71. message=r'^coroutine .* was never awaited$',
  72. category=RuntimeWarning)
  73. loop.call_soon_threadsafe(loop.stop)
  74. if __name__ == '__main__':
  75. sys.audit("cpython.run_stdin")
  76. loop = asyncio.new_event_loop()
  77. asyncio.set_event_loop(loop)
  78. repl_locals = {'asyncio': asyncio}
  79. for key in {'__name__', '__package__',
  80. '__loader__', '__spec__',
  81. '__builtins__', '__file__'}:
  82. repl_locals[key] = locals()[key]
  83. console = AsyncIOInteractiveConsole(repl_locals, loop)
  84. repl_future = None
  85. repl_future_interrupted = False
  86. try:
  87. import readline # NoQA
  88. except ImportError:
  89. pass
  90. repl_thread = REPLThread()
  91. repl_thread.daemon = True
  92. repl_thread.start()
  93. while True:
  94. try:
  95. loop.run_forever()
  96. except KeyboardInterrupt:
  97. if repl_future and not repl_future.done():
  98. repl_future.cancel()
  99. repl_future_interrupted = True
  100. continue
  101. else:
  102. break