Rewrite SocketManager

Author: Artucuno

fastapi_socketio/__init__.py

@@ -1 +1 @@
-from .socket_manager import SocketManager
+from .socket_manager import SocketManager

fastapi_socketio/socket_manager.py

@@ -2,8 +2,10 @@
 
 import socketio
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 
-class SocketManager:
+
+class SocketManager(socketio.AsyncServer):
     """
     Integrates SocketIO with FastAPI app.
     Adds `sio` property to FastAPI object (app).
@@ -18,88 +20,26 @@ class SocketManager:
     """
 
     def __init__(
-        self,
-        app: FastAPI,
-        mount_location: str = "/ws",
-        socketio_path: str = "socket.io",
-        cors_allowed_origins: Union[str, list] = '*',
-        async_mode: str = "asgi",
-        **kwargs
+            self,
+            app: FastAPI,
+            mount_location: str = "/ws",
+            socketio_path: str = "socket.io",
+            cors_allowed_origins: Union[str, list] = '*',
+            async_mode: str = "asgi",
+            **kwargs
     ) -> None:
-        # TODO: Change Cors policy based on fastapi cors Middleware
-        self._sio = socketio.AsyncServer(async_mode=async_mode, cors_allowed_origins=cors_allowed_origins, **kwargs)
+        middleware = next((x for x in app.user_middleware if issubclass(x.cls, CORSMiddleware)), None)
+        if middleware:
+            cors_allowed_origins = middleware.options.get("allow_origins", "*")
+        super().__init__(cors_allowed_origins=cors_allowed_origins, async_mode=async_mode, **kwargs)
         self._app = socketio.ASGIApp(
-            socketio_server=self._sio, socketio_path=socketio_path
+            socketio_server=self, socketio_path=socketio_path
         )
 
         app.mount(mount_location, self._app)
         app.add_route(f"/{socketio_path}/", route=self._app, methods=["GET", "POST"])
         app.add_websocket_route(f"/{socketio_path}/", self._app)
-        app.sio = self._sio
+        app.sio = self
 
     def is_asyncio_based(self) -> bool:
         return True
-
-    @property
-    def on(self):
-        return self._sio.on
-
-    @property
-    def attach(self):
-        return self._sio.attach
-
-    @property
-    def emit(self):
-        return self._sio.emit
-
-    @property
-    def send(self):
-        return self._sio.send
-
-    @property
-    def call(self):
-        return self._sio.call
-
-    @property
-    def close_room(self):
-        return self._sio.close_room
-
-    @property
-    def get_session(self):
-        return self._sio.get_session
-
-    @property
-    def save_session(self):
-        return self._sio.save_session
-
-    @property
-    def session(self):
-        return self._sio.session
-
-    @property
-    def disconnect(self):
-        return self._sio.disconnect
-
-    @property
-    def handle_request(self):
-        return self._sio.handle_request
-
-    @property
-    def start_background_task(self):
-        return self._sio.start_background_task
-
-    @property
-    def sleep(self):
-        return self._sio.sleep
-
-    @property
-    def enter_room(self):
-        return self._sio.enter_room
-
-    @property
-    def leave_room(self):
-        return self._sio.leave_room
-    
-    @property
-    def register_namespace(self):
-        return self._sio.register_namespace