Adds WebSocket API GW request

examples/rest/auth_example.py

@@ -1,4 +1,5 @@
 """Simple Lambda Handler with authorization"""
+
 from lbz.authz.decorators import authorization
 from lbz.dev.server import MyDevServer
 from lbz.exceptions import ServerError

lbz/authentication.py

@@ -1,4 +1,5 @@
 """JWT based Authentication module."""
+
 from lbz._cfg import AUTH_REMOVE_PREFIXES
 from lbz.jwt_utils import decode_jwt
 

lbz/dev/test.py

@@ -12,6 +12,7 @@ def __init__(self, resource: type[Resource]):
         self.resource = resource
 
     def post(
+        # pylint: disable=too-many-positional-arguments
         self,
         path: str,
         params: dict | None = None,
@@ -22,6 +23,7 @@ def post(
         return self._process(path, "POST", params, query_params, body, headers)
 
     def get(
+        # pylint: disable=too-many-positional-arguments
         self,
         path: str,
         params: dict | None = None,
@@ -32,6 +34,7 @@ def get(
         return self._process(path, "GET", params, query_params, body, headers)
 
     def patch(
+        # pylint: disable=too-many-positional-arguments
         self,
         path: str,
         params: dict | None = None,
@@ -42,6 +45,7 @@ def patch(
         return self._process(path, "PATCH", params, query_params, body, headers)
 
     def put(
+        # pylint: disable=too-many-positional-arguments
         self,
         path: str,
         params: dict | None = None,
@@ -52,6 +56,7 @@ def put(
         return self._process(path, "PUT", params, query_params, body, headers)
 
     def delete(
+        # pylint: disable=too-many-positional-arguments
         self,
         path: str,
         params: dict | None = None,
@@ -62,6 +67,7 @@ def delete(
         return self._process(path, "DELETE", params, query_params, body, headers)
 
     def _process(
+        # pylint: disable=too-many-positional-arguments
         self,
         path: str,
         method: str,

lbz/exceptions.py

@@ -340,7 +340,7 @@ class NetworkAuthenticationRequired(LambdaFWServerException):
 
 def all_lbz_errors(
     cls: type[LambdaFWException] = LambdaFWException,
-) -> Generator[type[LambdaFWException], None, None]:
+) -> Generator[type[LambdaFWException]]:
     for subcls in cls.__subclasses__():
         if subcls not in [LambdaFWClientException, LambdaFWServerException]:
             yield subcls

lbz/lambdas/client.py

@@ -56,6 +56,7 @@ def invoke(
 
     @classmethod
     def request(
+        # pylint: disable=too-many-positional-arguments
         cls,
         function_name: str,
         method: str,

lbz/misc.py

@@ -1,4 +1,5 @@
 """Misc Helpers of Lambda Framework."""
+
 from __future__ import annotations
 
 import copy

lbz/request.py

@@ -10,40 +10,22 @@
 from lbz.exceptions import BadRequestError
 from lbz.misc import MultiDict, get_logger
 from lbz.rest import ContentType
+from lbz.websocket import ActionType
 
 logger = get_logger(__name__)
 
 
-class Request:
-    """Represents request from API gateway."""
-
+class BaseRequest:
     def __init__(
         self,
-        headers: CIMultiDict,
-        uri_params: dict,
-        method: str,
         body: str | bytes | dict,
-        context: dict,
-        stage_vars: dict,
         is_base64_encoded: bool,
-        query_params: dict | None = None,
-        user: User | None = None,
     ):
-        self.query_params = MultiDict(query_params or {})
-        self.headers = headers
-        self.uri_params = uri_params
-        self.method = method
-        self.context = context
-        self.stage_vars = stage_vars
-        self.user = user
         self._is_base64_encoded = is_base64_encoded
         self._body = body
         self._json_body: dict | None = None
         self._raw_body: bytes | dict | None = None
 
-    def __repr__(self) -> str:
-        return f"<Request {self.method} >"
-
     @staticmethod
     def _decode_base64(encoded: str | bytes) -> bytes:
         if not isinstance(encoded, bytes):
@@ -68,6 +50,43 @@ def _safe_json_loads(payload: str | bytes) -> Any:
         except ValueError as error:
             raise BadRequestError(f"Invalid payload.\nPayload body:\n {repr(payload)}") from error
 
+
+class Request(BaseRequest):
+    """Represents request from HTTP API Gateway."""
+
+    def __init__(
+        # pylint: disable=too-many-positional-arguments
+        self,
+        headers: CIMultiDict,
+        uri_params: dict,
+        method: str,
+        body: str | bytes | dict,
+        context: dict,
+        stage_vars: dict,
+        is_base64_encoded: bool,
+        query_params: dict | None = None,
+        user: User | None = None,
+    ):
+        self.query_params = MultiDict(query_params or {})
+        self.headers = headers
+        self.uri_params = uri_params
+        self.method = method
+        self.context = context
+        self.stage_vars = stage_vars
+        self.user = user
+        super().__init__(body=body, is_base64_encoded=is_base64_encoded)
+
+    def __repr__(self) -> str:
+        return f"<Request {self.method} >"
+
+    def to_dict(self) -> dict:
+        copied = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
+        copied["headers"] = dict(copied["headers"])
+        copied["user"] = repr(copied["user"])
+        if copied["query_params"] is not None:
+            copied["query_params"] = dict(copied["query_params"])
+        return copied
+
     @property
     def json_body(self) -> dict | None:
         if self._json_body is None:
@@ -84,10 +103,48 @@ def json_body(self) -> dict | None:
                 raise BadRequestError(f"Content-Type header is missing or wrong: {content_type}")
         return self._json_body
 
-    def to_dict(self) -> dict:
-        copied = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
-        copied["headers"] = dict(copied["headers"])
-        copied["user"] = repr(copied["user"])
-        if copied["query_params"] is not None:
-            copied["query_params"] = dict(copied["query_params"])
-        return copied
+
+class WebSocketRequest(BaseRequest):
+    """Represents request from Web Socket Secure API Gateway."""
+
+    def __init__(
+        # pylint: disable=too-many-positional-arguments
+        self,
+        body: str | bytes | dict,
+        request_details: dict,
+        context: dict,
+        is_base64_encoded: bool,
+        user: User | None = None,
+        headers: CIMultiDict | None = None,
+    ):
+        self.headers = headers
+        self.context = context
+        self.user = user
+        self.action = request_details.pop("routeKey")
+        self.action_type = request_details.pop("eventType")
+        self.connection_id = request_details.pop("connectionId")
+        self.domain = request_details.pop("domainName")
+        self.stage = request_details.pop("stage")
+        self.details = request_details
+        super().__init__(body=body, is_base64_encoded=is_base64_encoded)
+
+    def __repr__(self) -> str:
+        return f"<Request {self.action_type} - {self.action} >"
+
+    def is_connection_request(self) -> bool:
+        return self.action_type is ActionType.CONNECT
+
+    def is_disconnection_request(self) -> bool:
+        return self.action_type is ActionType.DISCONNECT
+
+    @property
+    def json_body(self) -> dict | None:
+        if self._json_body is None:
+            if isinstance(self.raw_body, dict) or self.raw_body is None:
+                self._json_body = self.raw_body
+            else:
+                self._json_body = self._safe_json_loads(self.raw_body)
+        return self._json_body
+
+
+HTTPRequest = Request

lbz/rest/api_gateway_event.py

@@ -14,6 +14,7 @@ class APIGatewayEvent(dict):
     """
 
     def __init__(
+        # pylint: disable=too-many-positional-arguments
         self,
         method: str,
         resource_path: str,

lbz/type_defs.py

@@ -2,6 +2,7 @@
 
 https://docs.aws.amazon.com/lambda/latest/dg/python-context.html
 """
+
 from typing import Any
 
 

lbz/websocket/__init__.py

@@ -0,0 +1 @@
+from lbz.websocket.enums import ActionType

lbz/websocket/enums.py

@@ -0,0 +1,4 @@
+class ActionType:
+    CONNECT = "CONNECT"
+    DISCONNECT = "DISCONNECT"
+    MESSAGE = "MESSAGE"