Move to pyclass(frozen) (#232)

emmett-framework · Mar 15, 2024 · d8ecb12 · d8ecb12
1 parent 76ccb02
commit d8ecb12
Show file tree

Hide file tree

Showing 22 changed files with 790 additions and 682 deletions.
diff --git a/granian/_granian.pyi b/granian/_granian.pyi
@@ -4,31 +4,13 @@ from ._types import WebsocketMessage
 
 __version__: str
 
-class ASGIScope:
-    def as_dict(self, root_path: str) -> Dict[str, Any]: ...
-
 class RSGIHeaders:
     def __contains__(self, key: str) -> bool: ...
     def keys(self) -> List[str]: ...
     def values(self) -> List[str]: ...
     def items(self) -> List[Tuple[str]]: ...
     def get(self, key: str, default: Any = None) -> Any: ...
 
-class RSGIScope:
-    proto: str
-    http_version: str
-    rsgi_version: str
-    server: str
-    client: str
-    scheme: str
-    method: str
-    path: str
-    query_string: str
-    authority: Optional[str]
-
-    @property
-    def headers(self) -> RSGIHeaders: ...
-
 class RSGIHTTPStreamTransport:
     async def send_bytes(self, data: bytes): ...
     async def send_str(self, data: str): ...

diff --git a/granian/asgi.py b/granian/asgi.py
@@ -1,10 +1,15 @@
 import asyncio
 from functools import wraps
+from typing import Any, Dict
 
-from ._granian import ASGIScope as Scope
 from .log import logger
 
 
+class Scope:
+    def as_dict(self, root_path: str) -> Dict[str, Any]:
+        ...
+
+
 class LifespanProtocol:
     error_transition = 'Invalid lifespan state transition'
 

diff --git a/granian/rsgi.py b/granian/rsgi.py
@@ -1,16 +1,32 @@
 from enum import Enum
-from typing import Union
+from typing import Optional, Union
 
 from ._granian import (
-    RSGIHeaders as Headers,  # noqa
+    RSGIHeaders as Headers,
     RSGIHTTPProtocol as HTTPProtocol,  # noqa
     RSGIProtocolClosed as ProtocolClosed,  # noqa
     RSGIProtocolError as ProtocolError,  # noqa
-    RSGIScope as Scope,  # noqa
     RSGIWebsocketProtocol as WebsocketProtocol,  # noqa
 )
 
 
+class Scope:
+    proto: str
+    http_version: str
+    rsgi_version: str
+    server: str
+    client: str
+    scheme: str
+    method: str
+    path: str
+    query_string: str
+    authority: Optional[str]
+
+    @property
+    def headers(self) -> Headers:
+        ...
+
+
 class WebsocketMessageType(int, Enum):
     close = 0
     bytes = 1

diff --git a/src/asgi/callbacks.rs b/src/asgi/callbacks.rs
@@ -4,7 +4,8 @@ use tokio::sync::oneshot;
 
 use super::{
     io::{ASGIHTTPProtocol as HTTPProtocol, ASGIWebsocketProtocol as WebsocketProtocol, WebsocketDetachedTransport},
-    types::ASGIScope as Scope,
+    types::ASGIHTTPScope as HTTPScope,
+    types::ASGIWebsocketScope as WebsocketScope,
 };
 use crate::{
     callbacks::{
@@ -17,15 +18,15 @@ use crate::{
     ws::{HyperWebsocket, UpgradeData},
 };
 
-#[pyclass]
+#[pyclass(frozen)]
 pub(crate) struct CallbackRunnerHTTP {
     proto: Py<HTTPProtocol>,
     context: TaskLocals,
     cb: PyObject,
 }
 
 impl CallbackRunnerHTTP {
-    pub fn new(py: Python, cb: CallbackWrapper, proto: HTTPProtocol, scope: Scope) -> Self {
+    pub fn new(py: Python, cb: CallbackWrapper, proto: HTTPProtocol, scope: HTTPScope) -> Self {
         let pyproto = Py::new(py, proto).unwrap();
         Self {
             proto: pyproto.clone_ref(py),
@@ -51,23 +52,21 @@ impl CallbackRunnerHTTP {
 }
 
 macro_rules! callback_impl_done_http {
-    ($self:expr, $py:expr) => {
-        if let Ok(mut proto) = $self.proto.as_ref($py).try_borrow_mut() {
-            if let Some(tx) = proto.tx() {
-                let _ = tx.send(response_500());
-            }
+    ($self:expr) => {
+        if let Some(tx) = $self.proto.get().tx() {
+            let _ = tx.send(response_500());
         }
     };
 }
 
 macro_rules! callback_impl_done_err {
-    ($self:expr, $py:expr, $err:expr) => {
-        $self.done($py);
+    ($self:expr, $err:expr) => {
+        $self.done();
         log_application_callable_exception($err);
     };
 }
 
-#[pyclass]
+#[pyclass(frozen)]
 pub(crate) struct CallbackTaskHTTP {
     proto: Py<HTTPProtocol>,
     context: TaskLocals,
@@ -86,12 +85,12 @@ impl CallbackTaskHTTP {
         })
     }
 
-    fn done(&self, py: Python) {
-        callback_impl_done_http!(self, py);
+    fn done(&self) {
+        callback_impl_done_http!(self);
     }
 
-    fn err(&self, py: Python, err: &PyErr) {
-        callback_impl_done_err!(self, py, err);
+    fn err(&self, err: &PyErr) {
+        callback_impl_done_err!(self, err);
     }
 
     callback_impl_loop_run!();
@@ -109,7 +108,7 @@ impl CallbackTaskHTTP {
     }
 }
 
-#[pyclass]
+#[pyclass(frozen)]
 pub(crate) struct CallbackWrappedRunnerHTTP {
     #[pyo3(get)]
     proto: Py<HTTPProtocol>,
@@ -120,7 +119,7 @@ pub(crate) struct CallbackWrappedRunnerHTTP {
 }
 
 impl CallbackWrappedRunnerHTTP {
-    pub fn new(py: Python, cb: CallbackWrapper, proto: HTTPProtocol, scope: Scope) -> Self {
+    pub fn new(py: Python, cb: CallbackWrapper, proto: HTTPProtocol, scope: HTTPScope) -> Self {
         Self {
             proto: Py::new(py, proto).unwrap(),
             context: cb.context,
@@ -138,24 +137,24 @@ impl CallbackWrappedRunnerHTTP {
         callback_impl_loop_pytask!(pyself, py)
     }
 
-    fn done(&self, py: Python) {
-        callback_impl_done_http!(self, py);
+    fn done(&self) {
+        callback_impl_done_http!(self);
     }
 
-    fn err(&self, py: Python, err: &PyAny) {
-        callback_impl_done_err!(self, py, &PyErr::from_value(err));
+    fn err(&self, err: &PyAny) {
+        callback_impl_done_err!(self, &PyErr::from_value(err));
     }
 }
 
-#[pyclass]
+#[pyclass(frozen)]
 pub(crate) struct CallbackRunnerWebsocket {
     proto: Py<WebsocketProtocol>,
     context: TaskLocals,
     cb: PyObject,
 }
 
 impl CallbackRunnerWebsocket {
-    pub fn new(py: Python, cb: CallbackWrapper, proto: WebsocketProtocol, scope: Scope) -> Self {
+    pub fn new(py: Python, cb: CallbackWrapper, proto: WebsocketProtocol, scope: WebsocketScope) -> Self {
         let pyproto = Py::new(py, proto).unwrap();
         Self {
             proto: pyproto.clone(),
@@ -175,16 +174,14 @@ impl CallbackRunnerWebsocket {
 }
 
 macro_rules! callback_impl_done_ws {
-    ($self:expr, $py:expr) => {
-        if let Ok(mut proto) = $self.proto.as_ref($py).try_borrow_mut() {
-            if let (Some(tx), res) = proto.tx() {
-                let _ = tx.send(res);
-            }
+    ($self:expr) => {
+        if let (Some(tx), res) = $self.proto.get().tx() {
+            let _ = tx.send(res);
         }
     };
 }
 
-#[pyclass]
+#[pyclass(frozen)]
 pub(crate) struct CallbackTaskWebsocket {
     proto: Py<WebsocketProtocol>,
     context: TaskLocals,
@@ -203,12 +200,12 @@ impl CallbackTaskWebsocket {
         })
     }
 
-    fn done(&self, py: Python) {
-        callback_impl_done_ws!(self, py);
+    fn done(&self) {
+        callback_impl_done_ws!(self);
     }
 
-    fn err(&self, py: Python, err: &PyErr) {
-        callback_impl_done_err!(self, py, err);
+    fn err(&self, err: &PyErr) {
+        callback_impl_done_err!(self, err);
     }
 
     callback_impl_loop_run!();
@@ -226,7 +223,7 @@ impl CallbackTaskWebsocket {
     }
 }
 
-#[pyclass]
+#[pyclass(frozen)]
 pub(crate) struct CallbackWrappedRunnerWebsocket {
     #[pyo3(get)]
     proto: Py<WebsocketProtocol>,
@@ -237,7 +234,7 @@ pub(crate) struct CallbackWrappedRunnerWebsocket {
 }
 
 impl CallbackWrappedRunnerWebsocket {
-    pub fn new(py: Python, cb: CallbackWrapper, proto: WebsocketProtocol, scope: Scope) -> Self {
+    pub fn new(py: Python, cb: CallbackWrapper, proto: WebsocketProtocol, scope: WebsocketScope) -> Self {
         Self {
             proto: Py::new(py, proto).unwrap(),
             context: cb.context,
@@ -255,12 +252,12 @@ impl CallbackWrappedRunnerWebsocket {
         callback_impl_loop_pytask!(pyself, py)
     }
 
-    fn done(&self, py: Python) {
-        callback_impl_done_ws!(self, py);
+    fn done(&self) {
+        callback_impl_done_ws!(self);
     }
 
-    fn err(&self, py: Python, err: &PyAny) {
-        callback_impl_done_err!(self, py, &PyErr::from_value(err));
+    fn err(&self, err: &PyAny) {
+        callback_impl_done_err!(self, &PyErr::from_value(err));
     }
 }
 
@@ -293,7 +290,7 @@ macro_rules! call_impl_rtb_http {
             cb: CallbackWrapper,
             rt: RuntimeRef,
             req: HTTPRequest,
-            scope: Scope,
+            scope: HTTPScope,
         ) -> oneshot::Receiver<HTTPResponse> {
             let (tx, rx) = oneshot::channel();
             let protocol = HTTPProtocol::new(rt, req, tx);
@@ -313,7 +310,7 @@ macro_rules! call_impl_rtt_http {
             cb: CallbackWrapper,
             rt: RuntimeRef,
             req: HTTPRequest,
-            scope: Scope,
+            scope: HTTPScope,
         ) -> oneshot::Receiver<HTTPResponse> {
             let (tx, rx) = oneshot::channel();
             let protocol = HTTPProtocol::new(rt, req, tx);
@@ -336,7 +333,7 @@ macro_rules! call_impl_rtb_ws {
             rt: RuntimeRef,
             ws: HyperWebsocket,
             upgrade: UpgradeData,
-            scope: Scope,
+            scope: WebsocketScope,
         ) -> oneshot::Receiver<WebsocketDetachedTransport> {
             let (tx, rx) = oneshot::channel();
             let protocol = WebsocketProtocol::new(rt, tx, ws, upgrade);
@@ -357,7 +354,7 @@ macro_rules! call_impl_rtt_ws {
             rt: RuntimeRef,
             ws: HyperWebsocket,
             upgrade: UpgradeData,
-            scope: Scope,
+            scope: WebsocketScope,
         ) -> oneshot::Receiver<WebsocketDetachedTransport> {
             let (tx, rx) = oneshot::channel();
             let protocol = WebsocketProtocol::new(rt, tx, ws, upgrade);

diff --git a/src/asgi/http.rs b/src/asgi/http.rs
@@ -8,7 +8,8 @@ use super::{
         call_rtb_http, call_rtb_http_pyw, call_rtb_ws, call_rtb_ws_pyw, call_rtt_http, call_rtt_http_pyw, call_rtt_ws,
         call_rtt_ws_pyw,
     },
-    types::ASGIScope as Scope,
+    types::ASGIHTTPScope as HTTPScope,
+    types::ASGIWebsocketScope as WebsocketScope,
 };
 use crate::{
     callbacks::CallbackWrapper,
@@ -17,9 +18,9 @@ use crate::{
     ws::{is_upgrade_request as is_ws_upgrade, upgrade_intent as ws_upgrade, UpgradeData},
 };
 
-macro_rules! default_scope {
-    ($server_addr:expr, $client_addr:expr, $req:expr, $scheme:expr) => {
-        Scope::new(
+macro_rules! build_scope {
+    ($cls:ty, $server_addr:expr, $client_addr:expr, $req:expr, $scheme:expr) => {
+        <$cls>::new(
             $req.version(),
             $scheme,
             $req.uri().clone(),
@@ -53,7 +54,7 @@ macro_rules! handle_request {
             req: HTTPRequest,
             scheme: &str,
         ) -> HTTPResponse {
-            let scope = default_scope!(server_addr, client_addr, &req, scheme);
+            let scope = build_scope!(HTTPScope, server_addr, client_addr, &req, scheme);
             handle_http_response!($handler, rt, callback, req, scope)
         }
     };
@@ -69,10 +70,8 @@ macro_rules! handle_request_with_ws {
             req: HTTPRequest,
             scheme: &str,
         ) -> HTTPResponse {
-            let mut scope = default_scope!(server_addr, client_addr, &req, scheme);
-
             if is_ws_upgrade(&req) {
-                scope.set_websocket();
+                let scope = build_scope!(WebsocketScope, server_addr, client_addr, &req, scheme);
 
                 return match ws_upgrade(req, None) {
                     Ok((res, ws)) => {
@@ -130,6 +129,7 @@ macro_rules! handle_request_with_ws {
                 };
             }
 
+            let scope = build_scope!(HTTPScope, server_addr, client_addr, &req, scheme);
             handle_http_response!($handler_req, rt, callback, req, scope)
         }
     };