patch handle_exit in signal_shutdown instead of replacing signal hand…

…lers
AnswerDotAI · Jan 28, 2025 · 3f52bc9 · 3f52bc9
1 parent a2508e7
commit 3f52bc9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/fasthtml/core.py b/fasthtml/core.py
@@ -285,13 +285,13 @@ def EventStream(s):
 
 # %% ../nbs/api/00_core.ipynb
 def signal_shutdown():
+    from uvicorn.main import Server
     event = asyncio.Event()
-    def signal_handler(signum, frame):
+    @patch
+    def handle_exit(self:Server, *args, **kwargs):
         event.set()
-        signal.signal(signum, signal.SIG_DFL)
-        os.kill(os.getpid(), signum)
-
-    for sig in (signal.SIGINT, signal.SIGTERM): signal.signal(sig, signal_handler)
+        self.force_exit = True
+        self._orig_handle_exit(*args, **kwargs)
     return event
 
 # %% ../nbs/api/00_core.ipynb

diff --git a/nbs/api/00_core.ipynb b/nbs/api/00_core.ipynb
@@ -882,13 +882,13 @@
    "source": [
     "#| export\n",
     "def signal_shutdown():\n",
+    "    from uvicorn.main import Server\n",
     "    event = asyncio.Event()\n",
-    "    def signal_handler(signum, frame):\n",
+    "    @patch\n",
+    "    def handle_exit(self:Server, *args, **kwargs):\n",
     "        event.set()\n",
-    "        signal.signal(signum, signal.SIG_DFL)\n",
-    "        os.kill(os.getpid(), signum)\n",
-    "\n",
-    "    for sig in (signal.SIGINT, signal.SIGTERM): signal.signal(sig, signal_handler)\n",
+    "        self.force_exit = True\n",
+    "        self._orig_handle_exit(*args, **kwargs)\n",
     "    return event"
    ]
   },