From 3f812ec2b96b8ade782be63c533eb4c4e76ace30 Mon Sep 17 00:00:00 2001
From: CoolCat467 <52022020+CoolCat467@users.noreply.github.com>
Date: Sat, 2 Mar 2024 22:20:22 -0600
Subject: [PATCH] Test adding trio interopability

---
 src/statusbot/bot.py | 110 ++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 99 insertions(+), 11 deletions(-)

diff --git a/src/statusbot/bot.py b/src/statusbot/bot.py
index 4aaafbe..a432dc7 100644
--- a/src/statusbot/bot.py
+++ b/src/statusbot/bot.py
@@ -44,6 +44,8 @@
 
 import discord
 import mcstatus
+import outcome
+import trio
 
 # from discord.ext import tasks, commands
 from aiohttp.client_exceptions import ClientConnectorError
@@ -888,14 +890,17 @@ def __init__(
         self,
         prefix: str,
         loop: asyncio.AbstractEventLoop,
+        main_nursery: trio.Nursery,
+        trio_finish: trio.Event,
         *args: Any,
         intents: discord.Intents,
         **kwargs: Any,
     ) -> None:
         """Initialize StatusBot."""
         self.loop = loop
-        if "loop" in kwargs:
-            del kwargs["loop"]
+        self.nursery = main_nursery
+        self.trio_finish_event = trio_finish
+
         discord.Client.__init__(
             self,
             *args,
@@ -2101,6 +2106,7 @@ async def on_error(
     # Default, not affected by intents
     async def close(self) -> None:
         """Tell guilds bot shutting down."""
+        self.trio_finish_event.set()
         self.stopped.set()
         print("\nShutting down gears.")
         await gears.BaseBot.close(self)
@@ -2118,16 +2124,24 @@ async def close(self) -> None:
             await discord.Client.close(self)
 
 
-def run() -> None:
-    """Run bot."""
+def setup_bot(loop: asyncio.AbstractEventLoop) -> tuple[
+    tuple[
+        StatusBot,
+        asyncio.Task,
+    ],
+    tuple[
+        trio.Event,
+        asyncio.Future[outcome.Outcome],
+        asyncio.Future[trio.Nursery],
+    ],
+]:
+    """Return StatusBot run parts."""
     if TOKEN is None:
-        print(
+        raise RuntimeError(
             """\nNo token set!
 Either add ".env" file in bots folder with DISCORD_TOKEN=<token here> line,
 or set DISCORD_TOKEN environment variable.""",
         )
-        return
-    print("\nStarting bot...")
 
     intents = discord.Intents(
         dm_messages=True,
@@ -2140,15 +2154,89 @@ def run() -> None:
     )
     # 4867
 
+    trio_finish = trio.Event()
+
+    bot_run_task: asyncio.Task | None = None
+
+    trio_done_future = loop.create_future()
+
+    def trio_done_callback(trio_outcome: outcome.Outcome) -> None:
+        trio_done_future.set_result(trio_outcome)
+        print("Trio complete")
+        if (
+            isinstance(trio_outcome, outcome.Error)
+            and bot_run_task is not None
+        ):
+            print("[Trio] Canceling Bot Run")
+            # bot_run_task.set_exception(trio_outcome.error)
+            bot_run_task.cancel(
+                msg="".join(traceback.format_exception(trio_outcome.error)),
+            )
+        else:
+            trio_outcome.unwrap()
+
+    main_nursery: trio.Nursery | None = None
+
+    trio_nursery_future = loop.create_future()
+
+    @trio.lowlevel.disable_ki_protection
+    async def trio_async_root() -> None:
+        print("Trio start ticking.")
+        nonlocal main_nursery
+        async with trio.open_nursery() as nursery:
+            main_nursery = nursery
+            trio_nursery_future.set_result(nursery)
+            await trio_finish.wait()
+
+    trio.lowlevel.start_guest_run(
+        trio_async_root,
+        run_sync_soon_threadsafe=loop.call_soon_threadsafe,
+        run_sync_soon_not_threadsafe=loop.call_soon,
+        strict_exception_groups=True,
+        done_callback=trio_done_callback,
+    )
+
+    main_nursery = loop.run_until_complete(trio_nursery_future)
+
+    assert main_nursery is not None
+
+    bot = StatusBot(
+        BOT_PREFIX,
+        loop=loop,
+        intents=intents,
+        main_nursery=main_nursery,
+        trio_finish=trio_finish,
+    )
+
+    bot_run_task = loop.create_task(bot.start(TOKEN))
+    assert bot_run_task is not None
+
+    return (bot, bot_run_task), (trio_finish, trio_done_future, main_nursery)
+
+
+def run() -> None:
+    """Run bot."""
+    print("\nStarting bot...")
+
     loop = asyncio.new_event_loop()
-    bot = StatusBot(BOT_PREFIX, loop=loop, intents=intents)
+
+    (bot, bot_run_task), (trio_finish, trio_done_future, main_nursery) = (
+        setup_bot(loop)
+    )
 
     try:
-        loop.run_until_complete(bot.start(TOKEN))
+        loop.run_until_complete(bot_run_task)
     except KeyboardInterrupt:
-        print("\nClosing bot...")
-        loop.run_until_complete(bot.close())
+        print("Received KeyboardInterrupt")
+        trio_finish.set()
     finally:
+        print("\nShutting down bot...")
+        trio_finish.set()
+        loop.run_until_complete(bot.close())
+        # Cancel trio nursery
+        main_nursery.cancel_scope.cancel()
+        # Ensure closed
+        loop.run_until_complete(trio_done_future)
         # cancel all lingering tasks
         loop.close()
         print("\nBot has been deactivated.")