Fix: rewrite init method, override as_publish_kwargs

airtai · Jul 22, 2024 · 17d31b8 · 17d31b8
1 parent 7138709
commit 17d31b8
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 20 deletions.
diff --git a/faststream/confluent/response.py b/faststream/confluent/response.py
@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, Optional
 
+from typing_extensions import override
+
 from faststream.broker.response import Response
 
 if TYPE_CHECKING:
@@ -13,7 +15,25 @@ def __init__(
         *,
         headers: Optional["AnyDict"] = None,
         correlation_id: Optional[str] = None,
-        partition: Optional[int] = None,
         timestamp_ms: Optional[int] = None,
         key: Optional[bytes] = None,
-    ) -> None: ...
+    ) -> None:
+        super().__init__(
+            body=body,
+            headers=headers,
+            correlation_id=correlation_id,
+        )
+
+        self.timestamp_ms = timestamp_ms
+        self.key = key
+
+
+    @override
+    def as_publish_kwargs(self) -> "AnyDict":
+        publish_options = {
+            "headers": self.headers,
+            "correlation_id": self.correlation_id,
+            "timestamp_ms": self.timestamp_ms,
+            "key": self.key,
+        }
+        return publish_options
diff --git a/faststream/kafka/response.py b/faststream/kafka/response.py
@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, Optional
 
+from typing_extensions import override
+
 from faststream.broker.response import Response
 
 if TYPE_CHECKING:
@@ -13,7 +15,24 @@ def __init__(
         *,
         headers: Optional["AnyDict"] = None,
         correlation_id: Optional[str] = None,
-        partition: Optional[int] = None,
         timestamp_ms: Optional[int] = None,
         key: Optional[bytes] = None,
-    ) -> None: ...
+    ) -> None:
+        super().__init__(
+            body=body,
+            headers=headers,
+            correlation_id=correlation_id,
+        )
+
+        self.timestamp_ms = timestamp_ms
+        self.key = key
+
+    @override
+    def as_publish_kwargs(self) -> "AnyDict":
+        publish_options = {
+            "headers": self.headers,
+            "correlation_id": self.correlation_id,
+            "timestamp_ms": self.timestamp_ms,
+            "key": self.key,
+        }
+        return publish_options
diff --git a/faststream/nats/response.py b/faststream/nats/response.py
@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, Optional
 
+from typing_extensions import override
+
 from faststream.broker.response import Response
 
 if TYPE_CHECKING:
@@ -13,8 +15,17 @@ def __init__(
         *,
         headers: Optional["AnyDict"] = None,
         correlation_id: Optional[str] = None,
-        timeout: Optional[float] = None,
-        rpc: bool = False,
-        rpc_timeout: Optional[float] = 30.0,
-        raise_timeout: bool = False,
-    ) -> None: ...
+    ) -> None:
+        super().__init__(
+            body=body,
+            headers=headers,
+            correlation_id=correlation_id,
+        )
+
+    @override
+    def as_publish_kwargs(self) -> "AnyDict":
+        publish_options = {
+            "headers": self.headers,
+            "correlation_id": self.correlation_id,
+        }
+        return publish_options
diff --git a/faststream/rabbit/response.py b/faststream/rabbit/response.py
@@ -1,22 +1,65 @@
 from typing import TYPE_CHECKING, Optional
 
+from typing_extensions import override
+
 from faststream.broker.response import Response
 
 if TYPE_CHECKING:
-    from faststream.rabbit import RabbitQueue
+    from aio_pika.abc import DateType, TimeoutType
+
     from faststream.rabbit.types import AioPikaSendableMessage
+    from faststream.types import AnyDict
 
 
 class RabbitResponse(Response):
     def __init__(
         self,
         body: "AioPikaSendableMessage",
         *,
+        headers: Optional["AnyDict"] = None,
         correlation_id: Optional[str] = None,
-        queue: Optional["RabbitQueue"] = None,
-        routing_key: str = "",
         message_id: Optional[str] = None,
-        rpc: bool = False,
-        rpc_timeout: Optional[float] = 30.0,
-        raise_timeout: bool = False,
-    ) -> None: ...
+        mandatory: Optional[bool] = None,
+        immediate: Optional[bool] = None,
+        timeout: Optional["TimeoutType"] = None,
+        persist: Optional[bool] = None,
+        priority: Optional[int] = None,
+        message_type: Optional[str] = None,
+        content_type: Optional[str] = None,
+        expiration: Optional["DateType"] = None,
+        content_encoding: Optional[str] = None
+    ) -> None:
+        super().__init__(
+            body=body,
+            headers=headers,
+            correlation_id=correlation_id,
+        )
+
+        self.message_id = message_id
+        self.mandatory = mandatory
+        self.immediate = immediate
+        self.timeout = timeout
+        self.persist = persist
+        self.priority = priority
+        self.message_type = message_type
+        self.content_type = content_type
+        self.expiration = expiration
+        self.content_encoding = content_encoding
+
+    @override
+    def as_publish_kwargs(self) -> "AnyDict":
+        publish_options = {
+            "headers": self.headers,
+            "correlation_id": self.correlation_id,
+            "message_id": self.message_id,
+            "mandatory": self.mandatory,
+            "immediate": self.immediate,
+            "timeout": self.timeout,
+            "persist": self.persist,
+            "priority": self.priority,
+            "message_type": self.message_type,
+            "content_type": self.content_type,
+            "expiration": self.expiration,
+            "content_encoding": self.content_encoding
+        }
+        return publish_options
diff --git a/faststream/redis/response.py b/faststream/redis/response.py
@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, Optional
 
+from typing_extensions import override
+
 from faststream.broker.response import Response
 
 if TYPE_CHECKING:
@@ -13,7 +15,17 @@ def __init__(
         *,
         headers: Optional["AnyDict"] = None,
         correlation_id: Optional[str] = None,
-        rpc: bool = False,
-        rpc_timeout: Optional[float] = 30.0,
-        raise_timeout: bool = False,
-    ) -> None: ...
+    ) -> None:
+        super().__init__(
+            body=body,
+            headers=headers,
+            correlation_id=correlation_id,
+        )
+
+    @override
+    def as_publish_kwargs(self) -> "AnyDict":
+        publish_options = {
+            "headers": self.headers,
+            "correlation_id": self.correlation_id,
+        }
+        return publish_options