NatsBroker

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

    Args:
        exc_type: The type of the exception being handled, if any.
        exc_val: The exception instance being handled, if any.
        exec_tb: The traceback of the exception being handled, if any.

    Returns:
        None

    Raises:
        NotImplementedError: If the method is not implemented.
    !!! note

        The above docstring is autogenerated by docstring-gen library (https://docstring-gen.airt.ai)
    """
    super()._abc_close(exc_type, exc_val, exec_tb)

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

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

async def connect(
    self,
    *args: Any,
    **kwargs: Any,
) -> Client:
    connection = await super().connect(*args, **kwargs)
    for p in self._publishers.values():
        if p.stream is not None:
            p._producer = self._js_producer
        else:
            p._producer = self._producer
    return connection

def include_router(self, router: BrokerRouter[Any, MsgType]) -> None:
    """Includes a router in the current object.

    Args:
        router: The router to be included.

    Returns:
        None
    !!! note

        The above docstring is autogenerated by docstring-gen library (https://docstring-gen.airt.ai)
    """
    for r in router._handlers:
        self.subscriber(*r.args, **r.kwargs)(r.call)

    self._publishers.update(router._publishers)

def include_routers(self, *routers: BrokerRouter[Any, MsgType]) -> None:
    """Includes routers in the current object.

    Args:
        *routers: Variable length argument list of routers to include.

    Returns:
        None
    !!! note

        The above docstring is autogenerated by docstring-gen library (https://docstring-gen.airt.ai)
    """
    for r in routers:
        self.include_router(r)

@override
async def publish(  # type: ignore[override]
    self,
    *args: Any,
    stream: Optional[str] = None,
    **kwargs: Any,
) -> Optional[DecodedMessage]:
    if stream is None:
        assert self._producer, "NatsBroker is not started yet"  # nosec B101
        return await self._producer.publish(*args, **kwargs)
    else:
        assert self._js_producer, "NatsBroker is not started yet"  # nosec B101
        return await self._js_producer.publish(
            *args,
            stream=stream,
            **kwargs,  # type: ignore[misc]
        )

@override
def publisher(  # type: ignore[override]
    self,
    subject: str,
    headers: Optional[Dict[str, str]] = None,
    # Core
    reply_to: str = "",
    # JS
    stream: Union[str, JStream, None] = None,
    timeout: Optional[float] = None,
    # AsyncAPI information
    title: Optional[str] = None,
    description: Optional[str] = None,
    schema: Optional[Any] = None,
) -> Publisher:
    if (stream := stream_builder.stream(stream)) is not None:
        stream.subjects.append(subject)

    publisher = self._publishers.get(
        subject,
        Publisher(
            subject=subject,
            headers=headers,
            # Core
            reply_to=reply_to,
            # JS
            timeout=timeout,
            stream=stream,
            # AsyncAPI
            title=title,
            _description=description,
            _schema=schema,
        ),
    )
    super().publisher(subject, publisher)
    return publisher

async def start(self) -> None:
    context.set_local(
        "log_context",
        self._get_log_context(None, ""),
    )

    await super().start()
    assert (
        self._connection and self.stream
    ), "Broker should be started already"  # nosec B101

    for handler in self.handlers.values():
        stream = handler.stream

        if (is_js := stream is not None) and stream.declare:
            try:  # pragma: no branch
                await self.stream.add_stream(
                    config=stream.config,
                    subjects=stream.subjects,
                )

            except nats.js.errors.BadRequestError as e:
                old_config = (await self.stream.stream_info(stream.name)).config

                c = self._get_log_context(None, "")
                if (
                    e.description
                    == "stream name already in use with a different configuration"
                ):
                    self._log(str(e), logging.WARNING, c)
                    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, c, exc_info=e)

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

        c = self._get_log_context(
            None,
            subject=handler.subject,
            queue=handler.queue,
            stream=stream.name if stream else "",
        )
        self._log(f"`{handler.call_name}` waiting for messages", extra=c)
        await handler.start(self.stream if is_js else self._connection)

@override
def subscriber(  # type: ignore[override]
    self,
    subject: str,
    queue: str = "",
    pending_msgs_limit: Optional[int] = None,
    pending_bytes_limit: Optional[int] = None,
    # Core arguments
    max_msgs: int = 0,
    # JS arguments
    durable: Optional[str] = None,
    config: Optional[api.ConsumerConfig] = None,
    ordered_consumer: bool = False,
    idle_heartbeat: Optional[float] = None,
    flow_control: bool = False,
    deliver_policy: Optional[api.DeliverPolicy] = None,
    headers_only: Optional[bool] = None,
    # pull arguments
    pull_sub: Optional[PullSub] = None,
    inbox_prefix: bytes = api.INBOX_PREFIX,
    # custom
    ack_first: bool = False,
    stream: Union[str, JStream, None] = None,
    # broker arguments
    dependencies: Sequence[Depends] = (),
    parser: Optional[CustomParser[Msg, NatsMessage]] = None,
    decoder: Optional[CustomDecoder[NatsMessage]] = None,
    middlewares: Optional[Sequence[Callable[[Msg], BaseMiddleware]]] = None,
    filter: Filter[NatsMessage] = default_filter,
    # AsyncAPI information
    title: Optional[str] = None,
    description: Optional[str] = None,
    **original_kwargs: Any,
) -> Callable[
    [Callable[P_HandlerParams, T_HandlerReturn]],
    HandlerCallWrapper[Msg, P_HandlerParams, T_HandlerReturn],
]:
    stream = stream_builder.stream(stream)

    if pull_sub is not None and stream is None:
        raise ValueError("Pull subscriber can be used only with a stream")

    self._setup_log_context(
        queue=queue,
        subject=subject,
        stream=stream.name if stream else None,
    )
    super().subscriber()

    extra_options: AnyDict = {
        "pending_msgs_limit": pending_msgs_limit
        or (
            DEFAULT_JS_SUB_PENDING_MSGS_LIMIT
            if stream
            else DEFAULT_SUB_PENDING_MSGS_LIMIT
        ),
        "pending_bytes_limit": pending_bytes_limit
        or (
            DEFAULT_JS_SUB_PENDING_BYTES_LIMIT
            if stream
            else DEFAULT_SUB_PENDING_BYTES_LIMIT
        ),
    }

    if stream:
        extra_options.update(
            {
                "durable": durable,
                "stream": stream.name,
                "config": config,
            }
        )

        if pull_sub is not None:
            extra_options.update({"inbox_prefix": inbox_prefix})

        else:
            extra_options.update(
                {
                    "ordered_consumer": ordered_consumer,
                    "idle_heartbeat": idle_heartbeat,
                    "flow_control": flow_control,
                    "deliver_policy": deliver_policy,
                    "headers_only": headers_only,
                    "manual_ack": not ack_first,
                }
            )

    else:
        extra_options.update(
            {
                "max_msgs": max_msgs,
            }
        )

    key = Handler.get_routing_hash(subject)
    handler = self.handlers[key] = self.handlers.get(
        key,
        Handler(
            subject=subject,
            queue=queue,
            stream=stream,
            pull_sub=pull_sub,
            extra_options=extra_options,
            title=title,
            description=description,
            log_context_builder=partial(
                self._get_log_context,
                stream=stream.name if stream else "",
                subject=subject,
                queue=queue,
            ),
        ),
    )

    if stream:
        stream.subjects.append(handler.subject)

    def consumer_wrapper(
        func: Callable[P_HandlerParams, T_HandlerReturn],
    ) -> HandlerCallWrapper[Msg, P_HandlerParams, T_HandlerReturn,]:
        handler_call, dependant = self._wrap_handler(
            func,
            extra_dependencies=dependencies,
            **original_kwargs,
        )

        handler.add_call(
            handler=handler_call,
            filter=filter,
            middlewares=middlewares,
            parser=parser or self._global_parser,
            decoder=decoder or self._global_decoder,
            dependant=dependant,
        )

        return handler_call

    return consumer_wrapper