Modularise some code

robyn/__init__.py

@@ -35,8 +35,9 @@ def __init__(self, file_object):
         self.routes = []
         self.headers = []
         self.routes = []
-        self.directories = []
+        self.middlewares = []
         self.web_sockets = {}
+        self.directories = []
         self.event_handlers = {}
 
     def add_route(self, route_type, endpoint, handler):
@@ -61,6 +62,40 @@ def add_route(self, route_type, endpoint, handler):
             )
         )
 
+    def add_middleware_route(self, route_type, endpoint, handler):
+        """
+        [This is base handler for the middleware decorator]
+
+        :param route_type [str]: [??]
+        :param endpoint [str]: [endpoint for the route added]
+        :param handler [function]: [represents the sync or async function passed as a handler for the route]
+        """
+
+        """ We will add the status code here only
+        """
+        number_of_params = len(signature(handler).parameters)
+        self.routes.append(
+            (
+                route_type,
+                endpoint,
+                handler,
+                asyncio.iscoroutinefunction(handler),
+                number_of_params,
+            )
+        )
+
+    def before_request(self, endpoint):
+        """
+        [The @app.before_request decorator to add a get route]
+
+        :param endpoint [str]: [endpoint to server the route]
+        """
+
+        def inner(handler):
+            self.add_middleware_route("BEFORE_REQUEST", endpoint, handler)
+
+        return inner
+
     def add_directory(
         self, route, directory_path, index_file=None, show_files_listing=False
     ):

robyn/processpool.py

@@ -12,7 +12,7 @@
 
 
 def spawn_process(
-    directories, headers, routes, web_sockets, event_handlers, socket, workers
+    directories, headers, routes, middlewares, web_sockets, event_handlers, socket, workers
 ):
     """
     This function is called by the main process handler to create a server runtime.
@@ -21,6 +21,7 @@ def spawn_process(
     :param directories tuple: the list of all the directories and related data in a tuple
     :param headers tuple: All the global headers in a tuple
     :param routes tuple: The routes touple, containing the description about every route.
+    :param middlewares tuple: The middleware router touple, containing the description about every route.
     :param web_sockets list: This is a list of all the web socket routes
     :param event_handlers Dict: This is an event dict that contains the event handlers
     :param socket Socket: This is the main tcp socket, which is being shared across multiple processes.
@@ -53,6 +54,10 @@ def spawn_process(
         route_type, endpoint, handler, is_async, number_of_params = route
         server.add_route(route_type, endpoint, handler, is_async, number_of_params)
 
+    for route in middlewares:
+        route_type, endpoint, handler, is_async, number_of_params = route
+        server.add_middleware_route(route_type, endpoint, handler, is_async, number_of_params)
+
     if "startup" in event_handlers:
         server.add_startup_handler(event_handlers[Events.STARTUP][0], event_handlers[Events.STARTUP][1])
 

src/lib.rs

@@ -1,6 +1,5 @@
-mod middleware_router;
 mod processor;
-mod router;
+mod routers;
 mod server;
 mod shared_socket;
 mod types;

src/middleware_router.rs → src/routers/middleware_router.rs

@@ -1,5 +1,5 @@
 use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
+use std::sync::RwLock;
 // pyo3 modules
 use crate::types::PyFunction;
 use pyo3::prelude::*;
@@ -55,16 +55,12 @@ impl MiddlewareRouter {
 
     #[inline]
     fn get_relevant_map_str(&self, route: &str) -> Option<&RwLock<Node<(PyFunction, u8)>>> {
-        if route != "WS" {
-            let method = match Method::from_bytes(route.as_bytes()) {
-                Ok(res) => res,
-                Err(_) => return None,
-            };
+        let method = match Method::from_bytes(route.as_bytes()) {
+            Ok(res) => res,
+            Err(_) => return None,
+        };
 
-            return self.get_relevant_map(method);
-        } else {
-            return None;
-        }
+        return self.get_relevant_map(method);
     }
 
     // Checks if the functions is an async function

src/routers/mod.rs

@@ -0,0 +1,3 @@
+pub mod middleware_router;
+pub mod router;
+pub mod web_socket_router;

src/router.rs → src/routers/router.rs

@@ -1,5 +1,5 @@
 use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
+use std::sync::RwLock;
 // pyo3 modules
 use crate::types::PyFunction;
 use pyo3::prelude::*;

src/routers/web_socket_router.rs

@@ -0,0 +1,85 @@
+use std::collections::HashMap;
+use std::sync::RwLock;
+// pyo3 modules
+use crate::types::PyFunction;
+use pyo3::prelude::*;
+use pyo3::types::PyAny;
+
+use actix_web::http::Method;
+use matchit::Node;
+
+/// Contains the thread safe hashmaps of different routes
+
+pub struct WebSocketRouter {
+    web_socket_routes: RwLock<HashMap<String, HashMap<String, (PyFunction, u8)>>>,
+}
+
+impl WebSocketRouter {
+    pub fn new() -> Self {
+        Self {
+            web_socket_routes: RwLock::new(HashMap::new()),
+        }
+    }
+
+    #[inline]
+    pub fn get_web_socket_map(
+        &self,
+    ) -> &RwLock<HashMap<String, HashMap<String, (PyFunction, u8)>>> {
+        &self.web_socket_routes
+    }
+
+    // Checks if the functions is an async function
+    // Inserts them in the router according to their nature(CoRoutine/SyncFunction)
+    pub fn add_websocket_route(
+        &self,
+        route: &str,
+        connect_route: (Py<PyAny>, bool, u8),
+        close_route: (Py<PyAny>, bool, u8),
+        message_route: (Py<PyAny>, bool, u8),
+    ) {
+        let table = self.get_web_socket_map();
+        let (connect_route_function, connect_route_is_async, connect_route_params) = connect_route;
+        let (close_route_function, close_route_is_async, close_route_params) = close_route;
+        let (message_route_function, message_route_is_async, message_route_params) = message_route;
+
+        let insert_in_router =
+            |handler: Py<PyAny>, is_async: bool, number_of_params: u8, socket_type: &str| {
+                let function = if is_async {
+                    PyFunction::CoRoutine(handler)
+                } else {
+                    PyFunction::SyncFunction(handler)
+                };
+
+                println!("socket type is {:?} {:?}", table, route);
+
+                table
+                    .write()
+                    .unwrap()
+                    .entry(route.to_string())
+                    .or_default()
+                    .insert(socket_type.to_string(), (function, number_of_params))
+            };
+
+        insert_in_router(
+            connect_route_function,
+            connect_route_is_async,
+            connect_route_params,
+            "connect",
+        );
+
+        insert_in_router(
+            close_route_function,
+            close_route_is_async,
+            close_route_params,
+            "close",
+        );
+
+        insert_in_router(
+            message_route_function,
+            message_route_is_async,
+            message_route_params,
+            "message",
+        );
+    }
+
+}

src/server.rs

@@ -1,6 +1,6 @@
-use crate::middleware_router::MiddlewareRouter;
 use crate::processor::{apply_headers, execute_event_handler, handle_request};
-use crate::router::Router;
+use crate::routers::middleware_router::MiddlewareRouter;
+use crate::routers::router::Router;
 use crate::shared_socket::SocketHeld;
 use crate::types::{Headers, PyFunction};
 use crate::web_socket_connection::start_web_socket;
@@ -239,7 +239,8 @@ impl Server {
         number_of_params: u8,
     ) {
         println!("Route added for {} {} ", route_type, route);
-        self.middleware_router .add_route(route_type, route, handler, is_async, number_of_params);
+        self.middleware_router
+            .add_route(route_type, route, handler, is_async, number_of_params);
     }
 
     /// Add a new web socket route to the routing tables