added support for gateway controller inheritance

eadwinCode · eadwinCode · commit aab73fda48c0 · 2024-06-18T12:54:02.000+01:00
diff --git a/ellar/socket_io/decorators/gateway.py b/ellar/socket_io/decorators/gateway.py
@@ -10,7 +10,7 @@
     GATEWAY_WATERMARK,
 )
 from ellar.socket_io.model import GatewayBase, GatewayType
-from ellar.utils import get_name, get_type_of_base
+from ellar.utils import get_name
 
 
 def WebSocketGateway(
@@ -51,27 +51,15 @@ def _decorator(cls: t.Type) -> t.Type:
                 str(get_name(_gateway_type)).lower().replace("gateway", "")
             )
 
-        for base in get_type_of_base(GatewayBase, _gateway_type):
-            if reflect.has_metadata(GATEWAY_WATERMARK, base) and hasattr(
-                _gateway_type, "__GATEWAY_WATERMARK__"
-            ):
-                raise ImproperConfiguration(
-                    f"`@WebSocketGateway` decorated classes does not support inheritance. \n"
-                    f"{_gateway_type}"
-                )
-
-        if not reflect.has_metadata(GATEWAY_WATERMARK, _gateway_type) and not hasattr(
-            _gateway_type, "__GATEWAY_WATERMARK__"
-        ):
-            reflect.define_metadata(GATEWAY_WATERMARK, True, _gateway_type)
-            reflect.define_metadata(
-                GATEWAY_OPTIONS, _kwargs["socket_init_kwargs"], _gateway_type
-            )
+        reflect.define_metadata(GATEWAY_WATERMARK, True, _gateway_type)
+        reflect.define_metadata(
+            GATEWAY_OPTIONS, _kwargs["socket_init_kwargs"], _gateway_type
+        )
 
-            injectable(RequestORTransientScope)(_gateway_type)
+        injectable(RequestORTransientScope)(_gateway_type)
 
-            for key in GATEWAY_METADATA.keys:
-                reflect.define_metadata(key, _kwargs[key], _gateway_type)
+        for key in GATEWAY_METADATA.keys:
+            reflect.define_metadata(key, _kwargs[key], _gateway_type)
 
         if new_cls:
             _gateway_type.__GATEWAY_WATERMARK__ = True  # type:ignore[attr-defined]
diff --git a/ellar/socket_io/factory.py b/ellar/socket_io/factory.py
@@ -1,9 +1,7 @@
 import inspect
 import typing as t
-from abc import ABC
 
 import socketio
-from ellar.common.constants import CONTROLLER_CLASS_KEY
 from ellar.core.router_builders import RouterBuilder
 from ellar.reflect import reflect
 from ellar.socket_io.adapter import SocketIOASGIApp
@@ -30,7 +28,9 @@ def _get_message_handler(
         cls,
         klass: t.Type,
     ) -> t.Iterable[t.Union[t.Callable]]:
-        for method in klass.__dict__.values():
+        for _method_name, method in inspect.getmembers(
+            klass, predicate=inspect.isfunction
+        ):
             if hasattr(method, MESSAGE_MAPPING_METADATA) and getattr(
                 method, MESSAGE_MAPPING_METADATA
             ):
@@ -40,28 +40,16 @@ def _get_message_handler(
     def _process_controller_routes(
         cls, klass: t.Type[GatewayBase]
     ) -> t.List[t.Callable]:
-        bases = inspect.getmro(klass)
+        # bases = inspect.getmro(klass)
         results = []
 
         if reflect.get_metadata(GATEWAY_METADATA.PROCESSED, klass):
             return reflect.get_metadata(GATEWAY_MESSAGE_HANDLER_KEY, klass) or []
 
-        for base_cls in reversed(bases):
-            if base_cls not in [ABC, GatewayBase, object]:
-                for method in cls._get_message_handler(base_cls):
-                    if reflect.has_metadata(CONTROLLER_CLASS_KEY, method):
-                        raise Exception(
-                            f"{klass.__name__} Gateway message handler tried to be processed more than once."
-                            f"\n-Message Handler - {method}."
-                            f"\n-Gateway message handler can not be reused once its under a `@Gateway` decorator."
-                        )
+        for method in cls._get_message_handler(klass):
+            results.append(method)
 
-                    results.append(method)
-
-                    reflect.define_metadata(CONTROLLER_CLASS_KEY, klass, method)
-                    reflect.define_metadata(
-                        GATEWAY_MESSAGE_HANDLER_KEY, [method], klass
-                    )
+            reflect.define_metadata(GATEWAY_MESSAGE_HANDLER_KEY, [method], klass)
 
         reflect.define_metadata(GATEWAY_METADATA.PROCESSED, True, klass)
         return results
@@ -91,14 +79,18 @@ def build(
             is_disconnection_handler = reflect.get_metadata(DISCONNECT_EVENT, handler)
 
             if is_connection_handler:
-                SocketOperationConnection(CONNECTION_EVENT, socket_server, handler)
+                SocketOperationConnection(
+                    controller_type, CONNECTION_EVENT, socket_server, handler
+                )
             elif is_disconnection_handler:
-                SocketOperationConnection(DISCONNECT_EVENT, socket_server, handler)
+                SocketOperationConnection(
+                    controller_type, DISCONNECT_EVENT, socket_server, handler
+                )
             else:
                 message = reflect.get_metadata_or_raise_exception(
                     MESSAGE_METADATA, handler
                 )
-                SocketMessageOperation(message, socket_server, handler)
+                SocketMessageOperation(controller_type, message, socket_server, handler)
 
         return Mount(app=SocketIOASGIApp(socket_server), path=path, name=name)
 
diff --git a/ellar/socket_io/gateway.py b/ellar/socket_io/gateway.py
@@ -8,7 +8,6 @@
     serialize_object,
 )
 from ellar.common.constants import (
-    CONTROLLER_CLASS_KEY,
     CONTROLLER_OPERATION_HANDLER_KEY,
     EXTRA_ROUTE_ARGS_KEY,
     NOT_SET,
@@ -49,17 +48,19 @@ class SocketOperationConnection:
     )
 
     def __init__(
-        self, event: str, server: AsyncServer, message_handler: t.Callable
+        self,
+        controller_type: t.Type[GatewayBase],
+        event: str,
+        server: AsyncServer,
+        message_handler: t.Callable,
     ) -> None:
         self._event = event
         self._server = server
         self.endpoint = message_handler
         self._name = get_name(self.endpoint)
         self._is_coroutine = inspect.iscoroutinefunction(message_handler)
         self.endpoint_parameter_model = NOT_SET
-        self._controller_type: t.Type[GatewayBase] = reflect.get_metadata(  # type: ignore[assignment]
-            CONTROLLER_CLASS_KEY, self.endpoint
-        )
+        self._controller_type = controller_type
         self._load_model()
         self._register_handler()
         reflect.define_metadata(CONTROLLER_OPERATION_HANDLER_KEY, self, self.endpoint)
@@ -192,12 +193,6 @@ def get_controller_type(self) -> t.Type[GatewayBase]:
         For operation under ModuleRouter, this will return a unique type created for the router for tracking some properties
         :return: a type that wraps the operation
         """
-        if not self._controller_type:
-            _controller_type = reflect.get_metadata(CONTROLLER_CLASS_KEY, self.endpoint)
-            if _controller_type is None or not isinstance(_controller_type, type):
-                raise Exception("Operation must have a single control type.")
-            self._controller_type = t.cast(t.Type[GatewayBase], _controller_type)
-
         return self._controller_type
 
     def _get_gateway_instance(self, ctx: IExecutionContext) -> GatewayBase:
diff --git a/ellar/utils/__init__.py b/ellar/utils/__init__.py
@@ -76,7 +76,7 @@ def build_init_kwargs(obj: t.Type, init_kwargs: t.Dict) -> t.Dict:
 
 def get_type_of_base(
     base_type: t.Type[t.Any], reference_type: t.Type[t.Any]
-) -> t.Iterable[t.Type[t.Any]]:
+) -> t.Iterable[t.Type[t.Any]]:  # pragma: no cover
     for base in inspect.getmro(reference_type):
         if issubclass(base, base_type):
             yield base
diff --git a/tests/test_socket_io/test_decorators/test_gateway.py b/tests/test_socket_io/test_decorators/test_gateway.py
@@ -1,8 +1,7 @@
 import pytest
 from ellar.auth.guards import GuardHttpBearerAuth
 from ellar.common import UseGuards
-from ellar.common.constants import CONTROLLER_CLASS_KEY, GUARDS_KEY
-from ellar.common.exceptions import ImproperConfiguration
+from ellar.common.constants import CONTROLLER_OPERATION_HANDLER_KEY, GUARDS_KEY
 from ellar.reflect import reflect
 from ellar.socket_io import GatewayRouterFactory, WebSocketGateway, subscribe_message
 from ellar.socket_io.constants import (
@@ -97,28 +96,27 @@ def a_message(self):
     assert message == "a_message"
     assert (
         reflect.get_metadata_or_raise_exception(
-            CONTROLLER_CLASS_KEY, SampleAGateway().a_message
-        )
+            CONTROLLER_OPERATION_HANDLER_KEY, SampleAGateway().a_message
+        ).get_controller_type()
         is SampleAGateway
     )
 
 
-def test_sub_message_building_fails():
-    with pytest.raises(Exception) as ex:
-
-        @WebSocketGateway(path="/ws", namespace="/some-namespace")
-        class SampleBGateway(GatewayBase):
-            @subscribe_message
-            @reflect.metadata(CONTROLLER_CLASS_KEY, "b_message")
-            def b_message(self):
-                pass
-
-        GatewayRouterFactory.build(SampleBGateway)
-
-    assert (
-        "SampleBGateway Gateway message handler tried to be processed more than once"
-        in str(ex.value)
-    )
+# def test_sub_message_building_fails():
+#     with pytest.raises(Exception) as ex:
+#
+#         @WebSocketGateway(path="/ws", namespace="/some-namespace")
+#         class SampleBGateway(GatewayBase):
+#             @subscribe_message
+#             def b_message(self):
+#                 pass
+#
+#         GatewayRouterFactory.build(SampleBGateway)
+#
+#     assert (
+#         "SampleBGateway Gateway message handler tried to be processed more than once"
+#         in str(ex.value)
+#     )
 
 
 def test_cant_use_gateway_decorator_on_function():
@@ -131,13 +129,16 @@ def sample_c_gateway():
     assert "WebSocketGateway is a class decorator" in str(ex.value)
 
 
-def test_inheritance_fails():
-    with pytest.raises(ImproperConfiguration) as ex:
-
-        @WebSocketGateway(path="/ws", namespace="/some-namespace")
-        class InheritanceNotSupported(SampleWithoutGateway):
+def test_inheritance_works():
+    @WebSocketGateway(path="/ws", namespace="/some-namespace")
+    class InheritanceSupported(SampleWithoutGateway):
+        @subscribe_message
+        def b_message(self):
             pass
 
-    assert "`@WebSocketGateway` decorated classes does not support inheritance." in str(
-        ex.value
+    GatewayRouterFactory.build(InheritanceSupported)
+
+    message_handlers = reflect.get_metadata_or_raise_exception(
+        GATEWAY_MESSAGE_HANDLER_KEY, InheritanceSupported
     )
+    assert len(message_handlers) == 1
diff --git a/tests/test_socket_io/test_gateway.py b/tests/test_socket_io/test_gateway.py
