`patch_stdout` Causing Missing Print Statements on Application Exit

[Chris3606]

I noticed some odd behavior regarding patch_stdout when the application is exited, where print statements that are executed close to the application exiting aren't printed. To demonstrate, I modified the examples\prompts\async-prompt.py example to the following:

import asyncio

from prompt_toolkit.patch_stdout import patch_stdout
from prompt_toolkit.shortcuts import PromptSession

session = PromptSession("Say something: ")
do_exit = False
times_finally_called = 0
print_executed = False

async def print_counter(id: int, exit_on: int = -1):
    """
    Coroutine that prints counters.
    """
    global session, do_exit, times_finally_called

    try:
        i = 0
        while True:
            print("Counter %i: %i" % (id, i))
            if i == exit_on:
                session.app.exit()
                do_exit = True
                return

            i += 1
            await asyncio.sleep(1)
    except asyncio.CancelledError:
        print("Background task %i cancelled." % id)
    finally:
        print('Background task %i exiting.' % id)
        times_finally_called += 1

async def interactive_shell():
    """
    Like `interactive_shell`, but doing things manual.
    """
    global session, do_exit

    # Add background tasks
    session.app.create_background_task(print_counter(1))
    session.app.create_background_task(print_counter(2, 3))

    # Run echo loop.  I should probably use AppResult in exit and just use that to exit the loop
    # but this works for demonstration
    while not do_exit:
        try:
            result = await session.prompt_async()
            print('You said: "{0}"'.format(result))
        except (EOFError, KeyboardInterrupt):
            return

async def main():
    global print_executed
    with patch_stdout():
        await interactive_shell()
        print("Quitting event loop. Bye.")
        print_executed = True

if __name__ == "__main__":
    try:
        from asyncio import run
    except ImportError:
        print('using run until complete')
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main())
    else:
        print('using run')
        asyncio.run(main())

    print('Finally was actually called %i times' % times_finally_called)
    print('Print was executed: %s' % str(print_executed))

Basically, I wanted to force an application (prompt) exit from a coroutine. This example is a bit contrived, however it is less so in more practical applications like exiting the application as a result of an exception.

When running the above example in Python 3.7 and Python 3.8 on a Windows machine, the "Quitting event loop. Bye." message does not display, even though it was clearly executed because the variable assignment to True below the print statement did occur.

This appears to have something to do with how patch_stdout works or how it interacts with application exit, as simply moving that "Quitting event loop" message to outside of the with patch_stdout(): context causes it to display just fine.

I also tested this on Python 3.6, and the behavior is more complicated there. First, the original example was calling asyncio.run_until_complete, which I believe may be an error as that function doesn't exist -- run_until_complete appears to be only a function of event loops, and there is no module-level equivalent. Thus, I modified the code to simply create an event loop and call run_until_complete. However, the behavior here is quite a bit worse -- not only are most of the print statements involving the exit path missing, but the finally block in print_counter appears to be called only once instead of twice. I'm not quite sure what is occurring here, however I did notice that the run function added in 3.7 ensures that asynchronous generators are finalized (which run_until_complete does not), so I'd imagine that may have something to do with it.

[jonathanslenders]

That sounds like a bug. Thanks a lot for the extensive write-down! I really hope to look into this some time soon.

[Chris3606]

Thanks! I haven't had a lot of additional time to look into this myself, but feel free to let me know if there's anything I can do to help.

[nullus]

I ran into the same issue. Looking at patch_stdout/run_in_terminal, I've produced a minimal example which demonstrates what happens.

import asyncio

def run_in_terminal():
    print("run_in_terminal()")

    async def run():
        print("run()")

    return asyncio.ensure_future(run())

async def interactive_shell():
    print("interactive_shell()")

    # call_soon_threadsafe would be run from patch_stdout context handler
    asyncio.get_event_loop().call_soon_threadsafe(run_in_terminal)
    # This allows event loop to execute "run_in_terminal", so "run" 
    # will execute after "nothing":
    # await asyncio.sleep(0)
    await nothing()

async def nothing():
    print("nothing()")

if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(interactive_shell())
    # This will trigger run():
    # asyncio.get_event_loop().run_until_complete(nothing())

call_soon_threadsafe schedules run_in_terminal to run when the event loop resumes. When that runs ensure_future schedules run as a Task. On the final iteration of the event loop (after interactive_shell returns) it will run all the outstanding callbacks, but not tasks created by those callbacks.

The workaround is: either ensuring run_in_terminal runs and schedules run before leaving interactive_shell (e.g. await asyncio.sleep(0)), or starting the event loop again to process them.

I'm looking to put together a PR, but I'm still figuring out the right approach. In the example here, simply replacing self.loop.call_soon_threadsafe(write_and_flush_in_loop) with write_and_flush_in_loop() in StdoutProxy has the desired effect, though I'm not certain that would work for all use cases.

[Chris3606]

@nullus Thank you for taking the time to minimize the example. I did play around with removing the call_soon_threadsafe call as you had explained, and it did appear to solve the issue for this example (and a more complex one in which I had initially discovered the issue). However, to be honest I don't understand why (and this might just be me lacking some understanding of asyncio). As far as I can tell, the ensure_future(run()) code ultimately schedules a Task, not a callback -- so I don't understand how it is guaranteed to complete when the event loop finishes, when compared to calling via call_soon_threadsafe?

[Chris3606]

@jonathanslenders Would you accept a PR for this that simply adds a flag to patch_stdout which, when set, omits the call_soon_threadsafe call and just calls write_and_flush_in_loop directly? Obviously only works for situations where you can guarantee that you're in the same thread as the event loop, but I have no idea how else to fix this.

[Chris3606]

Actually, I don't think this workaround even works in modern versions anymore.