__main__.py 3.3 KB

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