codekingpro/portable-devtools

1import argparse2import ast3import asyncio4import asyncio.tools5import concurrent.futures6import contextvars7import inspect8import os9import site10import sys11import threading12import types13import warnings14 15from _colorize import get_theme16from _pyrepl.console import InteractiveColoredConsole17 18from . import futures19 20 21class AsyncIOInteractiveConsole(InteractiveColoredConsole):22 23    def __init__(self, locals, loop):24        super().__init__(locals, filename="<stdin>")25        self.compile.compiler.flags |= ast.PyCF_ALLOW_TOP_LEVEL_AWAIT26 27        self.loop = loop28        self.context = contextvars.copy_context()29 30    def runcode(self, code):31        global return_code32        future = concurrent.futures.Future()33 34        def callback():35            global return_code36            global repl_future37            global keyboard_interrupted38 39            repl_future = None40            keyboard_interrupted = False41 42            func = types.FunctionType(code, self.locals)43            try:44                coro = func()45            except SystemExit as se:46                return_code = se.code47                self.loop.stop()48                return49            except KeyboardInterrupt as ex:50                keyboard_interrupted = True51                future.set_exception(ex)52                return53            except BaseException as ex:54                future.set_exception(ex)55                return56 57            if not inspect.iscoroutine(coro):58                future.set_result(coro)59                return60 61            try:62                repl_future = self.loop.create_task(coro, context=self.context)63                futures._chain_future(repl_future, future)64            except BaseException as exc:65                future.set_exception(exc)66 67        self.loop.call_soon_threadsafe(callback, context=self.context)68 69        try:70            return future.result()71        except SystemExit as se:72            return_code = se.code73            self.loop.stop()74            return75        except BaseException:76            if keyboard_interrupted:77                if not CAN_USE_PYREPL:78                    self.write("\nKeyboardInterrupt\n")79            else:80                self.showtraceback()81            return self.STATEMENT_FAILED82 83class REPLThread(threading.Thread):84 85    def run(self):86        global return_code87 88        try:89            if not sys.flags.quiet:90                banner = (91                    f'asyncio REPL {sys.version} on {sys.platform}\n'92                    f'Use "await" directly instead of "asyncio.run()".\n'93                    f'Type "help", "copyright", "credits" or "license" '94                    f'for more information.\n'95                )96 97                console.write(banner)98 99            if not sys.flags.isolated and (startup_path := os.getenv("PYTHONSTARTUP")):100                sys.audit("cpython.run_startup", startup_path)101 102                import tokenize103                with tokenize.open(startup_path) as f:104                    startup_code = compile(f.read(), startup_path, "exec")105                    exec(startup_code, console.locals)106 107            ps1 = getattr(sys, "ps1", ">>> ")108            if CAN_USE_PYREPL:109                theme = get_theme().syntax110                ps1 = f"{theme.prompt}{ps1}{theme.reset}"111            console.write(f"{ps1}import asyncio\n")112 113            if CAN_USE_PYREPL:114                from _pyrepl.simple_interact import (115                    run_multiline_interactive_console,116                )117                try:118                    run_multiline_interactive_console(console)119                except SystemExit:120                    # expected via the `exit` and `quit` commands121                    pass122                except BaseException:123                    # unexpected issue124                    console.showtraceback()125                    console.write("Internal error, ")126                    return_code = 1127            else:128                console.interact(banner="", exitmsg="")129        finally:130            warnings.filterwarnings(131                'ignore',132                message=r'^coroutine .* was never awaited$',133                category=RuntimeWarning)134 135            loop.call_soon_threadsafe(loop.stop)136 137    def interrupt(self) -> None:138        if not CAN_USE_PYREPL:139            return140 141        from _pyrepl.simple_interact import _get_reader142        r = _get_reader()143        if r.threading_hook is not None:144            r.threading_hook.add("")  # type: ignore145 146 147if __name__ == '__main__':148    parser = argparse.ArgumentParser(149        prog="python3 -m asyncio",150        description="Interactive asyncio shell and CLI tools",151        color=True,152    )153    subparsers = parser.add_subparsers(help="sub-commands", dest="command")154    ps = subparsers.add_parser(155        "ps", help="Display a table of all pending tasks in a process"156    )157    ps.add_argument("pid", type=int, help="Process ID to inspect")158    pstree = subparsers.add_parser(159        "pstree", help="Display a tree of all pending tasks in a process"160    )161    pstree.add_argument("pid", type=int, help="Process ID to inspect")162    args = parser.parse_args()163    match args.command:164        case "ps":165            asyncio.tools.display_awaited_by_tasks_table(args.pid)166            sys.exit(0)167        case "pstree":168            asyncio.tools.display_awaited_by_tasks_tree(args.pid)169            sys.exit(0)170        case None:171            pass  # continue to the interactive shell172        case _:173            # shouldn't happen as an invalid command-line wouldn't parse174            # but let's keep it for the next person adding a command175            print(f"error: unhandled command {args.command}", file=sys.stderr)176            parser.print_usage(file=sys.stderr)177            sys.exit(1)178 179    sys.audit("cpython.run_stdin")180 181    if os.getenv('PYTHON_BASIC_REPL'):182        CAN_USE_PYREPL = False183    else:184        from _pyrepl.main import CAN_USE_PYREPL185 186    return_code = 0187    loop = asyncio.new_event_loop()188    asyncio.set_event_loop(loop)189 190    repl_locals = {'asyncio': asyncio}191    for key in {'__name__', '__package__',192                '__loader__', '__spec__',193                '__builtins__', '__file__'}:194        repl_locals[key] = locals()[key]195 196    console = AsyncIOInteractiveConsole(repl_locals, loop)197 198    repl_future = None199    keyboard_interrupted = False200 201    try:202        import readline  # NoQA203    except ImportError:204        readline = None205 206    interactive_hook = getattr(sys, "__interactivehook__", None)207 208    if interactive_hook is not None:209        sys.audit("cpython.run_interactivehook", interactive_hook)210        interactive_hook()211 212    if interactive_hook is site.register_readline:213        # Fix the completer function to use the interactive console locals214        try:215            import rlcompleter216        except:217            pass218        else:219            if readline is not None:220                completer = rlcompleter.Completer(console.locals)221                readline.set_completer(completer.complete)222 223    repl_thread = REPLThread(name="Interactive thread")224    repl_thread.daemon = True225    repl_thread.start()226 227    while True:228        try:229            loop.run_forever()230        except KeyboardInterrupt:231            keyboard_interrupted = True232            if repl_future and not repl_future.done():233                repl_future.cancel()234            repl_thread.interrupt()235            continue236        else:237            break238 239    console.write('exiting asyncio REPL...\n')240    loop.close()241    sys.exit(return_code)242