chore: polish types

airtai · Dec 15, 2024 · d958655 · d958655
1 parent d1a2376
commit d958655
Show file tree

Hide file tree

Showing 11 changed files with 135 additions and 119 deletions.
diff --git a/faststream/kafka/broker/broker.py b/faststream/kafka/broker/broker.py
@@ -667,49 +667,47 @@ def _subscriber_setup_extra(self) -> "AnyDict":
 
     @overload
     async def publish(
-            self,
-            message: "SendableMessage",
-            topic: str = "",
-            *,
-            key: Union[bytes, Any, None] = None,
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            correlation_id: Optional[str] = None,
-            reply_to: str = "",
-            no_confirm: Literal[True] = False,
-    ) -> "asyncio.Future[RecordMetadata]":
-        ...
+        self,
+        message: "SendableMessage",
+        topic: str = "",
+        *,
+        key: Union[bytes, Any, None] = None,
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        correlation_id: Optional[str] = None,
+        reply_to: str = "",
+        no_confirm: Literal[True],
+    ) -> "asyncio.Future[RecordMetadata]": ...
 
     @overload
     async def publish(
-            self,
-            message: "SendableMessage",
-            topic: str = "",
-            *,
-            key: Union[bytes, Any, None] = None,
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            correlation_id: Optional[str] = None,
-            reply_to: str = "",
-            no_confirm: Literal[False] = False,
-    ) -> "RecordMetadata":
-        ...
+        self,
+        message: "SendableMessage",
+        topic: str = "",
+        *,
+        key: Union[bytes, Any, None] = None,
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        correlation_id: Optional[str] = None,
+        reply_to: str = "",
+        no_confirm: Literal[False] = False,
+    ) -> "RecordMetadata": ...
 
     @override
     async def publish(
-            self,
-            message: "SendableMessage",
-            topic: str = "",
-            *,
-            key: Union[bytes, Any, None] = None,
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            correlation_id: Optional[str] = None,
-            reply_to: str = "",
-            no_confirm: bool = False,
+        self,
+        message: "SendableMessage",
+        topic: str = "",
+        *,
+        key: Union[bytes, Any, None] = None,
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        correlation_id: Optional[str] = None,
+        reply_to: str = "",
+        no_confirm: bool = False,
     ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]:
         """Publish message directly.
 
@@ -842,42 +840,44 @@ async def request(  # type: ignore[override]
 
     @overload
     async def publish_batch(
-            self,
-            *messages: "SendableMessage",
-            topic: str = "",
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            reply_to: str = "",
-            correlation_id: Optional[str] = None,
-            no_confirm: Literal[True] = False,
+        self,
+        *messages: "SendableMessage",
+        topic: str = "",
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        reply_to: str = "",
+        correlation_id: Optional[str] = None,
+        no_confirm: Literal[True],
     ) -> "asyncio.Future[RecordMetadata]": ...
 
     @overload
     async def publish_batch(
-            self,
-            *messages: "SendableMessage",
-            topic: str = "",
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            reply_to: str = "",
-            correlation_id: Optional[str] = None,
-            no_confirm: Literal[False] = False,
+        self,
+        *messages: "SendableMessage",
+        topic: str = "",
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        reply_to: str = "",
+        correlation_id: Optional[str] = None,
+        no_confirm: Literal[False] = False,
     ) -> "RecordMetadata": ...
 
     async def publish_batch(
-            self,
-            *messages: "SendableMessage",
-            topic: str = "",
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            reply_to: str = "",
-            correlation_id: Optional[str] = None,
-            no_confirm: bool = False,
+        self,
+        *messages: "SendableMessage",
+        topic: str = "",
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        reply_to: str = "",
+        correlation_id: Optional[str] = None,
+        no_confirm: bool = False,
     ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]:
-        """Args:
+        """Publish a message batch as a single request to broker.
+
+        Args:
             *messages:
                 Messages bodies to send.
             topic:

diff --git a/faststream/kafka/publisher/producer.py b/faststream/kafka/publisher/producer.py
@@ -48,7 +48,7 @@ async def disconnect(self) -> None:
         await self._producer.stop()
         self._producer = EmptyProducerState()
 
-    def __bool__(self) -> None:
+    def __bool__(self) -> bool:
         return bool(self._producer)
 
     @property

diff --git a/faststream/kafka/publisher/state.py b/faststream/kafka/publisher/state.py
@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from typing import TYPE_CHECKING, Protocol
 
 from faststream.exceptions import IncorrectState
@@ -8,7 +9,10 @@
 
 class ProducerState(Protocol):
     producer: "AIOKafkaProducer"
-    closed: bool
+
+    @property
+    @abstractmethod
+    def closed(self) -> bool: ...
 
     def __bool__(self) -> bool: ...
 

diff --git a/faststream/kafka/publisher/usecase.py b/faststream/kafka/publisher/usecase.py
@@ -155,32 +155,32 @@ def __init__(
 
     @overload
     async def publish(
-            self,
-            message: "SendableMessage",
-            topic: str = "",
-            *,
-            key: Union[bytes, Any, None] = None,
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            correlation_id: Optional[str] = None,
-            reply_to: str = "",
-            no_confirm: Literal[True] = False,
+        self,
+        message: "SendableMessage",
+        topic: str = "",
+        *,
+        key: Union[bytes, Any, None] = None,
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        correlation_id: Optional[str] = None,
+        reply_to: str = "",
+        no_confirm: Literal[True],
     ) -> "asyncio.Future[RecordMetadata]": ...
 
     @overload
     async def publish(
-            self,
-            message: "SendableMessage",
-            topic: str = "",
-            *,
-            key: Union[bytes, Any, None] = None,
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            correlation_id: Optional[str] = None,
-            reply_to: str = "",
-            no_confirm: Literal[False] = False,
+        self,
+        message: "SendableMessage",
+        topic: str = "",
+        *,
+        key: Union[bytes, Any, None] = None,
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        correlation_id: Optional[str] = None,
+        reply_to: str = "",
+        no_confirm: Literal[False] = False,
     ) -> "RecordMetadata": ...
 
     @override
@@ -197,7 +197,9 @@ async def publish(
         reply_to: str = "",
         no_confirm: bool = False,
     ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]:
-        """Args:
+        """Publishes a message to Kafka.
+
+        Args:
             message:
                 Message body to send.
             topic:
@@ -336,31 +338,30 @@ async def request(
 
 
 class BatchPublisher(LogicPublisher[tuple["ConsumerRecord", ...]]):
-
     @overload
     async def publish(
-            self,
-            *messages: "SendableMessage",
-            topic: str = "",
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            reply_to: str = "",
-            correlation_id: Optional[str] = None,
-            no_confirm: Literal[True] = False,
+        self,
+        *messages: "SendableMessage",
+        topic: str = "",
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        reply_to: str = "",
+        correlation_id: Optional[str] = None,
+        no_confirm: Literal[True],
     ) -> "asyncio.Future[RecordMetadata]": ...
 
     @overload
     async def publish(
-            self,
-            *messages: "SendableMessage",
-            topic: str = "",
-            partition: Optional[int] = None,
-            timestamp_ms: Optional[int] = None,
-            headers: Optional[dict[str, str]] = None,
-            reply_to: str = "",
-            correlation_id: Optional[str] = None,
-            no_confirm: Literal[False] = False,
+        self,
+        *messages: "SendableMessage",
+        topic: str = "",
+        partition: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        headers: Optional[dict[str, str]] = None,
+        reply_to: str = "",
+        correlation_id: Optional[str] = None,
+        no_confirm: Literal[False] = False,
     ) -> "RecordMetadata": ...
 
     @override
@@ -375,7 +376,9 @@ async def publish(
         correlation_id: Optional[str] = None,
         no_confirm: bool = False,
     ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]:
-        """Args:
+        """Publish a message batch as a single request to broker.
+
+        Args:
             *messages:
                 Messages bodies to send.
             topic:

diff --git a/faststream/specification/asyncapi/message.py b/faststream/specification/asyncapi/message.py
@@ -18,11 +18,11 @@
 
 def parse_handler_params(call: "CallModel", prefix: str = "") -> AnyDict:
     """Parses the handler parameters."""
-    model = getattr(call, "serializer", call).model  # type: ignore[union-attr]
+    model = getattr(call, "serializer", call).model
     assert model  # nosec B101
 
     body = get_model_schema(
-        create_model(  # type: ignore[call-overload]
+        create_model(
             model.__name__,
             **{p.field_name: (p.field_type, p.default_value) for p in call.flat_params},
         ),

diff --git a/faststream/specification/asyncapi/v2_6_0/facade.py b/faststream/specification/asyncapi/v2_6_0/facade.py
@@ -61,7 +61,7 @@ def to_yaml(self) -> str:
         return self.schema.to_yaml()
 
     @property
-    def schema(self) -> ApplicationSchema:  # type: ignore[override]
+    def schema(self) -> ApplicationSchema:
         return get_app_schema(
             self.broker,
             title=self.title,

diff --git a/faststream/specification/asyncapi/v3_0_0/facade.py b/faststream/specification/asyncapi/v3_0_0/facade.py
@@ -61,7 +61,7 @@ def to_yaml(self) -> str:
         return self.schema.to_yaml()
 
     @property
-    def schema(self) -> ApplicationSchema:  # type: ignore[override]
+    def schema(self) -> ApplicationSchema:
         return get_app_schema(
             self.broker,
             title=self.title,

diff --git a/faststream/specification/asyncapi/v3_0_0/schema/__init__.py b/faststream/specification/asyncapi/v3_0_0/schema/__init__.py
@@ -1,5 +1,3 @@
-from faststream.specification.asyncapi.v2_6_0.schema import ServerVariable
-
 from .channels import Channel
 from .components import Components
 from .contact import Contact
@@ -9,7 +7,7 @@
 from .message import CorrelationId, Message
 from .operations import Operation
 from .schema import ApplicationSchema
-from .servers import Server
+from .servers import Server, ServerVariable
 from .tag import Tag
 from .utils import Parameter, Reference
 

diff --git a/faststream/specification/asyncapi/v3_0_0/schema/servers.py b/faststream/specification/asyncapi/v3_0_0/schema/servers.py
@@ -10,6 +10,12 @@
 SecurityRequirement = list[dict[str, list[str]]]
 
 
+__all__ = (
+    "Server",
+    "ServerVariable",
+)
+
+
 class Server(BaseModel):
     """A class to represent a server.
 

diff --git a/faststream/specification/base/specification.py b/faststream/specification/base/specification.py
@@ -1,11 +1,14 @@
+from abc import abstractmethod
 from typing import Any, Protocol, runtime_checkable
 
 from .schema import BaseApplicationSchema
 
 
 @runtime_checkable
 class Specification(Protocol):
-    schema: BaseApplicationSchema
+    @property
+    @abstractmethod
+    def schema(self) -> BaseApplicationSchema: ...
 
     def to_json(self) -> str:
         return self.schema.to_json()

diff --git a/tests/brokers/kafka/test_publish.py b/tests/brokers/kafka/test_publish.py
@@ -151,7 +151,9 @@ async def handler(m) -> None:
         async with self.patch_broker(pub_broker) as br:
             await br.start()
 
-            batch_record_metadata_future = await br.publish_batch(1, "hi", topic=queue, no_confirm=True)
+            batch_record_metadata_future = await br.publish_batch(
+                1, "hi", topic=queue, no_confirm=True
+            )
             record_metadata_future = await br.publish("", topic=queue, no_confirm=True)
             assert isinstance(batch_record_metadata_future, asyncio.Future)
             assert isinstance(record_metadata_future, asyncio.Future)