From 2afd3c15b649ea9da92bbbc5c132637a455a5588 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Pierre?= <sebastien.pierre@gmail.com>
Date: Wed, 2 Oct 2024 13:14:28 +1300
Subject: [PATCH] [New] http,server: onClose response support

---
 src/py/extra/__init__.py   | 7 ++++++-
 src/py/extra/http/model.py | 9 ++++++++-
 src/py/extra/server.py     | 6 ++++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/py/extra/__init__.py b/src/py/extra/__init__.py
index 5195b7e..4093aa0 100644
--- a/src/py/extra/__init__.py
+++ b/src/py/extra/__init__.py
@@ -1,4 +1,9 @@
-from .http.model import HTTPRequest, HTTPResponse, HTTPRequestError  # NOQA: F401
+from .http.model import (
+    HTTPRequest,
+    HTTPResponse,
+    HTTPResponseLine,
+    HTTPRequestError,
+)  # NOQA: F401
 from .decorators import on, expose, pre, post  # NOQA: F401
 from .server import run  # NOQA: F401
 from .model import Service  # NOQA: F401
diff --git a/src/py/extra/http/model.py b/src/py/extra/http/model.py
index b5e5691..b703a84 100644
--- a/src/py/extra/http/model.py
+++ b/src/py/extra/http/model.py
@@ -551,7 +551,7 @@ def Create(
             protocol=protocol,
         )
 
-    __slots__ = ["protocol", "status", "message", "headers", "body"]
+    __slots__ = ["protocol", "status", "message", "headers", "body", "_onClose"]
 
     def __init__(
         self,
@@ -569,6 +569,7 @@ def __init__(
         # don't support that.
         self.headers: HTTPHeaders = headers
         self.body: THTTPBody | None = body
+        self._onClose: Callable[[HTTPResponse], None] | None = None
 
     # TODO: Deprecate
     def getHeader(self, name: str) -> str | None:
@@ -600,6 +601,12 @@ def head(self) -> bytes:
         # TODO: UTF8 maybe? Why ASCII?
         return "\r\n".join(lines).encode("ascii")
 
+    def onClose(
+        self, callback: Callable[["HTTPResponse"], None] | None
+    ) -> "HTTPResponse":
+        self._onClose = callback
+        return self
+
     def __str__(self) -> str:
         return f"Response({self.protocol} {self.status} {self.message} {self.headers} {self.body})"
 
diff --git a/src/py/extra/server.py b/src/py/extra/server.py
index b6df65d..783f11e 100644
--- a/src/py/extra/server.py
+++ b/src/py/extra/server.py
@@ -288,6 +288,12 @@ async def SendResponse(
             except Exception as e:
                 # NOTE: close handler failed
                 exception(e)
+        if res and res._onClose:
+            try:
+                res._onClose(res)
+            except Exception as e:
+                # NOTE: close handler failed
+                exception(e)
         if req and not sent:
             try:
                 warning(