NatsBroker

def setup(self) -> None:
    """Prepare all Broker entities to startup."""
    for h in self._subscribers.values():
        self.setup_subscriber(h)

    for p in self._publishers.values():
        self.setup_publisher(p)

def add_middleware(self, middleware: "BrokerMiddleware[MsgType]") -> None:
    """Append BrokerMiddleware to the end of middlewares list.

    Current middleware will be used as a most inner of already existed ones.
    """
    self._middlewares = (*self._middlewares, middleware)

    for sub in self._subscribers.values():
        sub.add_middleware(middleware)

    for pub in self._publishers.values():
        pub.add_middleware(middleware)

@override
def subscriber(  # type: ignore[override]
    self,
    subject: Annotated[
        str,
        Doc("NATS subject to subscribe."),
    ] = "",
    queue: Annotated[
        str,
        Doc(
            "Subscribers' NATS queue name. Subscribers with same queue name will be load balanced by the NATS "
            "server."
        ),
    ] = "",
    pending_msgs_limit: Annotated[
        Optional[int],
        Doc(
            "Limit of messages, considered by NATS server as possible to be delivered to the client without "
            "been answered. In case of NATS Core, if that limits exceeds, you will receive NATS 'Slow Consumer' "
            "error. "
            "That's literally means that your worker can't handle the whole load. In case of NATS JetStream, "
            "you will no longer receive messages until some of delivered messages will be acked in any way."
        ),
    ] = None,
    pending_bytes_limit: Annotated[
        Optional[int],
        Doc(
            "The number of bytes, considered by NATS server as possible to be delivered to the client without "
            "been answered. In case of NATS Core, if that limit exceeds, you will receive NATS 'Slow Consumer' "
            "error."
            "That's literally means that your worker can't handle the whole load. In case of NATS JetStream, "
            "you will no longer receive messages until some of delivered messages will be acked in any way."
        ),
    ] = None,
    # Core arguments
    max_msgs: Annotated[
        int,
        Doc("Consuming messages limiter. Automatically disconnect if reached."),
    ] = 0,
    # JS arguments
    durable: Annotated[
        Optional[str],
        Doc(
            "Name of the durable consumer to which the the subscription should be bound."
        ),
    ] = None,
    config: Annotated[
        Optional["api.ConsumerConfig"],
        Doc("Configuration of JetStream consumer to be subscribed with."),
    ] = None,
    ordered_consumer: Annotated[
        bool,
        Doc("Enable ordered consumer mode."),
    ] = False,
    idle_heartbeat: Annotated[
        Optional[float],
        Doc("Enable Heartbeats for a consumer to detect failures."),
    ] = None,
    flow_control: Annotated[
        bool,
        Doc("Enable Flow Control for a consumer."),
    ] = False,
    deliver_policy: Annotated[
        Optional["api.DeliverPolicy"],
        Doc("Deliver Policy to be used for subscription."),
    ] = None,
    headers_only: Annotated[
        Optional[bool],
        Doc(
            "Should be message delivered without payload, only headers and metadata."
        ),
    ] = None,
    # pull arguments
    pull_sub: Annotated[
        Union[bool, "PullSub"],
        Doc(
            "NATS Pull consumer parameters container. "
            "Should be used with `stream` only."
        ),
    ] = False,
    kv_watch: Annotated[
        Union[str, "KvWatch", None],
        Doc("KeyValue watch parameters container."),
    ] = None,
    obj_watch: Annotated[
        Union[bool, "ObjWatch"],
        Doc("ObjecStore watch parameters container."),
    ] = False,
    inbox_prefix: Annotated[
        bytes,
        Doc(
            "Prefix for generating unique inboxes, subjects with that prefix and NUID."
        ),
    ] = api.INBOX_PREFIX,
    # custom
    ack_first: Annotated[
        bool,
        Doc("Whether to `ack` message at start of consuming or not."),
    ] = False,
    stream: Annotated[
        Union[str, "JStream", None],
        Doc("Subscribe to NATS Stream with `subject` filter."),
    ] = None,
    # broker arguments
    dependencies: Annotated[
        Iterable["Depends"],
        Doc("Dependencies list (`[Depends(),]`) to apply to the subscriber."),
    ] = (),
    parser: Annotated[
        Optional["CustomCallable"],
        Doc("Parser to map original **nats-py** Msg to FastStream one."),
    ] = None,
    decoder: Annotated[
        Optional["CustomCallable"],
        Doc("Function to decode FastStream msg bytes body to python objects."),
    ] = None,
    middlewares: Annotated[
        Iterable["SubscriberMiddleware[NatsMessage]"],
        Doc("Subscriber middlewares to wrap incoming message processing."),
    ] = (),
    filter: Annotated[
        Union[
            "Filter[NatsMessage]",
            "Filter[NatsBatchMessage]",
        ],
        Doc(
            "Overload subscriber to consume various messages from the same source."
        ),
        deprecated(
            "Deprecated in **FastStream 0.5.0**. "
            "Please, create `subscriber` object and use it explicitly instead. "
            "Argument will be removed in **FastStream 0.6.0**."
        ),
    ] = default_filter,
    max_workers: Annotated[
        int,
        Doc("Number of workers to process messages concurrently."),
    ] = 1,
    retry: Annotated[
        bool,
        Doc("Whether to `nack` message at processing exception."),
    ] = False,
    no_ack: Annotated[
        bool,
        Doc("Whether to disable **FastStream** autoacknowledgement logic or not."),
    ] = False,
    no_reply: Annotated[
        bool,
        Doc(
            "Whether to disable **FastStream** RPC and Reply To auto responses or not."
        ),
    ] = False,
    # AsyncAPI information
    title: Annotated[
        Optional[str],
        Doc("AsyncAPI subscriber object title."),
    ] = None,
    description: Annotated[
        Optional[str],
        Doc(
            "AsyncAPI subscriber object description. "
            "Uses decorated docstring as default."
        ),
    ] = None,
    include_in_schema: Annotated[
        bool,
        Doc("Whetever to include operation in AsyncAPI schema or not."),
    ] = True,
) -> AsyncAPISubscriber:
    """Creates NATS subscriber object.

    You can use it as a handler decorator `@broker.subscriber(...)`.
    """
    stream = self._stream_builder.create(stream)

    subscriber = cast(
        AsyncAPISubscriber,
        super().subscriber(
            create_subscriber(
                subject=subject,
                queue=queue,
                stream=stream,
                pull_sub=PullSub.validate(pull_sub),
                kv_watch=KvWatch.validate(kv_watch),
                obj_watch=ObjWatch.validate(obj_watch),
                max_workers=max_workers,
                # extra args
                pending_msgs_limit=pending_msgs_limit,
                pending_bytes_limit=pending_bytes_limit,
                max_msgs=max_msgs,
                durable=durable,
                config=config,
                ordered_consumer=ordered_consumer,
                idle_heartbeat=idle_heartbeat,
                flow_control=flow_control,
                deliver_policy=deliver_policy,
                headers_only=headers_only,
                inbox_prefix=inbox_prefix,
                ack_first=ack_first,
                # subscriber args
                no_ack=no_ack,
                no_reply=no_reply,
                retry=retry,
                broker_middlewares=self._middlewares,
                broker_dependencies=self._dependencies,
                # AsyncAPI
                title_=title,
                description_=description,
                include_in_schema=self._solve_include_in_schema(include_in_schema),
            )
        ),
    )

    if stream and subscriber.subject:
        stream.add_subject(subscriber.subject)

    return subscriber.add_call(
        filter_=filter,
        parser_=parser or self._parser,
        decoder_=decoder or self._decoder,
        dependencies_=dependencies,
        middlewares_=middlewares,
    )

@override
def publisher(  # type: ignore[override]
    self,
    subject: Annotated[
        str,
        Doc("NATS subject to send message."),
    ],
    *,
    headers: Annotated[
        Optional[Dict[str, str]],
        Doc(
            "Message headers to store metainformation. "
            "**content-type** and **correlation_id** will be set automatically by framework anyway. "
            "Can be overridden by `publish.headers` if specified."
        ),
    ] = None,
    reply_to: Annotated[
        str,
        Doc("NATS subject name to send response."),
    ] = "",
    # JS
    stream: Annotated[
        Union[str, "JStream", None],
        Doc(
            "This option validates that the target `subject` is in presented stream. "
            "Can be omitted without any effect."
        ),
    ] = None,
    timeout: Annotated[
        Optional[float],
        Doc("Timeout to send message to NATS."),
    ] = None,
    # basic args
    middlewares: Annotated[
        Iterable["PublisherMiddleware"],
        Doc("Publisher middlewares to wrap outgoing messages."),
    ] = (),
    # AsyncAPI information
    title: Annotated[
        Optional[str],
        Doc("AsyncAPI publisher object title."),
    ] = None,
    description: Annotated[
        Optional[str],
        Doc("AsyncAPI publisher object description."),
    ] = None,
    schema: Annotated[
        Optional[Any],
        Doc(
            "AsyncAPI publishing message type. "
            "Should be any python-native object annotation or `pydantic.BaseModel`."
        ),
    ] = None,
    include_in_schema: Annotated[
        bool,
        Doc("Whetever to include operation in AsyncAPI schema or not."),
    ] = True,
) -> "AsyncAPIPublisher":
    """Creates long-living and AsyncAPI-documented publisher object.

    You can use it as a handler decorator (handler should be decorated by `@broker.subscriber(...)` too) - `@broker.publisher(...)`.
    In such case publisher will publish your handler return value.

    Or you can create a publisher object to call it lately - `broker.publisher(...).publish(...)`.
    """
    stream = self._stream_builder.create(stream)

    publisher = cast(
        AsyncAPIPublisher,
        super().publisher(
            publisher=AsyncAPIPublisher.create(
                subject=subject,
                headers=headers,
                # Core
                reply_to=reply_to,
                # JS
                timeout=timeout,
                stream=stream,
                # Specific
                broker_middlewares=self._middlewares,
                middlewares=middlewares,
                # AsyncAPI
                title_=title,
                description_=description,
                schema_=schema,
                include_in_schema=self._solve_include_in_schema(include_in_schema),
            )
        ),
    )

    if stream and publisher.subject:
        stream.add_subject(publisher.subject)

    return publisher

@override
def include_router(  # type: ignore[override]
    self,
    router: "NatsRegistrator",
    *,
    prefix: str = "",
    dependencies: Iterable["Depends"] = (),
    middlewares: Iterable["BrokerMiddleware[Msg]"] = (),
    include_in_schema: Optional[bool] = None,
) -> None:
    sub_streams = router._stream_builder.objects.copy()

    sub_router_subjects = [sub.subject for sub in router._subscribers.values()]

    for stream in sub_streams.values():
        new_subjects = []
        for subj in stream.subjects:
            if subj in sub_router_subjects:
                new_subjects.append("".join((self.prefix, subj)))
            else:
                new_subjects.append(subj)
        stream.subjects = new_subjects

    self._stream_builder.objects.update(sub_streams)

    return super().include_router(
        router,
        prefix=prefix,
        dependencies=dependencies,
        middlewares=middlewares,
        include_in_schema=include_in_schema,
    )

def include_routers(
    self,
    *routers: "ABCBroker[MsgType]",
) -> None:
    """Includes routers in the object."""
    for r in routers:
        self.include_router(r)

def get_fmt(self) -> str:
    """Fallback method to get log format if `log_fmt` if not specified."""
    return (
        "%(asctime)s %(levelname)-8s - "
        + (f"%(stream)-{self._max_stream_len}s | " if self._max_stream_len else "")
        + (f"%(queue)-{self._max_queue_len}s | " if self._max_queue_len else "")
        + f"%(subject)-{self._max_subject_len}s | "
        + f"%(message_id)-{self.__max_msg_id_ln}s - "
        "%(message)s"
    )

async def close(
    self,
    exc_type: Optional[Type[BaseException]] = None,
    exc_val: Optional[BaseException] = None,
    exc_tb: Optional["TracebackType"] = None,
) -> None:
    """Closes the object."""
    self.running = False

    for h in self._subscribers.values():
        await h.close()

    if self._connection is not None:
        await self._close(exc_type, exc_val, exc_tb)

@override
async def connect(  # type: ignore[override]
    self,
    servers: Annotated[
        Union[str, Iterable[str]],
        Doc("NATS cluster addresses to connect."),
    ] = EMPTY,
    **kwargs: "Unpack[NatsInitKwargs]",
) -> "Client":
    """Connect broker object to NATS cluster.

    To startup subscribers too you should use `broker.start()` after/instead this method.
    """
    if servers is not EMPTY:
        connect_kwargs: AnyDict = {
            **kwargs,
            "servers": servers,
        }
    else:
        connect_kwargs = {**kwargs}

    return await super().connect(**connect_kwargs)

async def start(self) -> None:
    """Connect broker to NATS cluster and startup all subscribers."""
    await super().start()

    assert self._connection  # nosec B101
    assert self.stream, "Broker should be started already"  # nosec B101
    assert self._producer, "Broker should be started already"  # nosec B101

    for stream in filter(
        lambda x: x.declare,
        self._stream_builder.objects.values(),
    ):
        try:
            await self.stream.add_stream(
                config=stream.config,
                subjects=stream.subjects,
            )

        except BadRequestError as e:  # noqa: PERF203
            log_context = AsyncAPISubscriber.build_log_context(
                message=None,
                subject="",
                queue="",
                stream=stream.name,
            )

            if (
                e.description
                == "stream name already in use with a different configuration"
            ):
                old_config = (await self.stream.stream_info(stream.name)).config

                self._log(str(e), logging.WARNING, log_context)
                await self.stream.update_stream(
                    config=stream.config,
                    subjects=tuple(
                        set(old_config.subjects or ()).union(stream.subjects)
                    ),
                )

            else:  # pragma: no cover
                self._log(str(e), logging.ERROR, log_context, exc_info=e)

        finally:
            # prevent from double declaration
            stream.declare = False

    # TODO: filter by already running handlers after TestClient refactor
    for handler in self._subscribers.values():
        self._log(
            f"`{handler.call_name}` waiting for messages",
            extra=handler.get_log_context(None),
        )
        await handler.start()

@override
async def publish(  # type: ignore[override]
    self,
    message: Annotated[
        "SendableMessage",
        Doc(
            "Message body to send. "
            "Can be any encodable object (native python types or `pydantic.BaseModel`)."
        ),
    ],
    subject: Annotated[
        str,
        Doc("NATS subject to send message."),
    ],
    headers: Annotated[
        Optional[Dict[str, str]],
        Doc(
            "Message headers to store metainformation. "
            "**content-type** and **correlation_id** will be set automatically by framework anyway."
        ),
    ] = None,
    reply_to: Annotated[
        str,
        Doc("NATS subject name to send response."),
    ] = "",
    correlation_id: Annotated[
        Optional[str],
        Doc(
            "Manual message **correlation_id** setter. "
            "**correlation_id** is a useful option to trace messages."
        ),
    ] = None,
    stream: Annotated[
        Optional[str],
        Doc(
            "This option validates that the target subject is in presented stream. "
            "Can be omitted without any effect."
        ),
    ] = None,
    timeout: Annotated[
        Optional[float],
        Doc("Timeout to send message to NATS."),
    ] = None,
    *,
    rpc: Annotated[
        bool,
        Doc("Whether to wait for reply in blocking mode."),
        deprecated(
            "Deprecated in **FastStream 0.5.17**. "
            "Please, use `request` method instead. "
            "Argument will be removed in **FastStream 0.6.0**."
        ),
    ] = False,
    rpc_timeout: Annotated[
        Optional[float],
        Doc("RPC reply waiting time."),
        deprecated(
            "Deprecated in **FastStream 0.5.17**. "
            "Please, use `request` method with `timeout` instead. "
            "Argument will be removed in **FastStream 0.6.0**."
        ),
    ] = 30.0,
    raise_timeout: Annotated[
        bool,
        Doc(
            "Whetever to raise `TimeoutError` or return `None` at **rpc_timeout**. "
            "RPC request returns `None` at timeout by default."
        ),
        deprecated(
            "Deprecated in **FastStream 0.5.17**. "
            "`request` always raises TimeoutError instead. "
            "Argument will be removed in **FastStream 0.6.0**."
        ),
    ] = False,
) -> Optional["DecodedMessage"]:
    """Publish message directly.

    This method allows you to publish message in not AsyncAPI-documented way. You can use it in another frameworks
    applications or to publish messages from time to time.

    Please, use `@broker.publisher(...)` or `broker.publisher(...).publish(...)` instead in a regular way.
    """
    publish_kwargs = {
        "subject": subject,
        "headers": headers,
        "reply_to": reply_to,
        "rpc": rpc,
        "rpc_timeout": rpc_timeout,
        "raise_timeout": raise_timeout,
    }

    producer: Optional[ProducerProto]
    if stream is None:
        producer = self._producer
    else:
        producer = self._js_producer
        publish_kwargs.update(
            {
                "stream": stream,
                "timeout": timeout,
            }
        )

    return await super().publish(
        message,
        producer=producer,
        correlation_id=correlation_id or gen_cor_id(),
        **publish_kwargs,
    )

@override
async def request(  # type: ignore[override]
    self,
    message: Annotated[
        "SendableMessage",
        Doc(
            "Message body to send. "
            "Can be any encodable object (native python types or `pydantic.BaseModel`)."
        ),
    ],
    subject: Annotated[
        str,
        Doc("NATS subject to send message."),
    ],
    headers: Annotated[
        Optional[Dict[str, str]],
        Doc(
            "Message headers to store metainformation. "
            "**content-type** and **correlation_id** will be set automatically by framework anyway."
        ),
    ] = None,
    correlation_id: Annotated[
        Optional[str],
        Doc(
            "Manual message **correlation_id** setter. "
            "**correlation_id** is a useful option to trace messages."
        ),
    ] = None,
    stream: Annotated[
        Optional[str],
        Doc(
            "This option validates that the target subject is in presented stream. "
            "Can be omitted without any effect."
        ),
    ] = None,
    timeout: Annotated[
        float,
        Doc("Timeout to send message to NATS."),
    ] = 0.5,
) -> "NatsMessage":
    publish_kwargs = {
        "subject": subject,
        "headers": headers,
        "timeout": timeout,
    }

    producer: Optional[ProducerProto]
    if stream is None:
        producer = self._producer

    else:
        producer = self._js_producer
        publish_kwargs.update({"stream": stream})

    msg: NatsMessage = await super().request(
        message,
        producer=producer,
        correlation_id=correlation_id or gen_cor_id(),
        **publish_kwargs,
    )
    return msg

@override
def setup_subscriber(  # type: ignore[override]
    self,
    subscriber: "AsyncAPISubscriber",
) -> None:
    connection: Union[
        Client,
        JetStreamContext,
        KVBucketDeclarer,
        OSBucketDeclarer,
        None,
    ] = None

    if getattr(subscriber, "kv_watch", None):
        connection = self._kv_declarer

    elif getattr(subscriber, "obj_watch", None):
        connection = self._os_declarer

    elif getattr(subscriber, "stream", None):
        connection = self.stream

    else:
        connection = self._connection

    return super().setup_subscriber(
        subscriber,
        connection=connection,
    )

@override
def setup_publisher(  # type: ignore[override]
    self,
    publisher: "AsyncAPIPublisher",
) -> None:
    producer: Optional[ProducerProto] = None

    if publisher.stream is not None:
        if self._js_producer is not None:
            producer = self._js_producer

    elif self._producer is not None:
        producer = self._producer

    super().setup_publisher(publisher, producer=producer)

async def key_value(
    self,
    bucket: str,
    *,
    description: Optional[str] = None,
    max_value_size: Optional[int] = None,
    history: int = 1,
    ttl: Optional[float] = None,  # in seconds
    max_bytes: Optional[int] = None,
    storage: Optional["StorageType"] = None,
    replicas: int = 1,
    placement: Optional["Placement"] = None,
    republish: Optional["RePublish"] = None,
    direct: Optional[bool] = None,
    # custom
    declare: bool = True,
) -> "KeyValue":
    assert self._kv_declarer, "Broker should be connected already."  # nosec B101

    return await self._kv_declarer.create_key_value(
        bucket=bucket,
        description=description,
        max_value_size=max_value_size,
        history=history,
        ttl=ttl,
        max_bytes=max_bytes,
        storage=storage,
        replicas=replicas,
        placement=placement,
        republish=republish,
        direct=direct,
        declare=declare,
    )