_namespaces.py

__all__ = ("namespace_context", "namespaced_id")

from contextlib import contextmanager
from contextvars import ContextVar, Token
import functools

from typing import Awaitable, Union, Optional, TypeVar, Callable, cast

from shiny.types import MISSING, MISSING_TYPE

from ._utils import is_async_callable

_current_namespace: ContextVar[Optional[str]] = ContextVar(
    "current_namespace", default=None
)


def get_current_namespace() -> Optional[str]:
    """
    Get the current namespace.
    """
    return _current_namespace.get()


@contextmanager
def namespace_context(ns: str):
    """
    Set a namespace for the duration of the context.
    """
    token: Token[Union[str, None]] = _current_namespace.set(ns)
    try:
        yield
    finally:
        _current_namespace.reset(token)


def namespaced_id(id: str, ns: Union[str, MISSING_TYPE, None] = MISSING) -> str:
    """
    Namespace an ID based on the current ``Module()``'s namespace.

    Parameters
    ----------
    id
        The ID to namespace.

    Warning
    -------
    This is only provided so that htmltools can use it to namespace ids within a
    ``Module()`` UI function.
    """
    if isinstance(ns, MISSING_TYPE):
        ns = _current_namespace.get()

    if ns is None:
        return id
    else:
        return ns + "_" + id


T = TypeVar("T")


def add_namespace(
    fn: Union[Callable[..., T], Callable[..., Awaitable[T]]],
    ns: Union[str, MISSING_TYPE, None] = MISSING,
) -> Union[Callable[..., T], Callable[..., Awaitable[T]]]:
    """
    Modify a function to use the **current** namespace when it executes.

    Note
    ----
    This is useful for marking user supplied functions to use the **current** namespace
    when they execute **later on**.
    """

    if isinstance(ns, MISSING_TYPE):
        ns = _current_namespace.get()

    if ns is None:
        return fn

    if is_async_callable(fn):

        @functools.wraps(fn)
        async def wrapper(*args: object, **kwargs: object) -> T:
            with namespace_context(ns):
                return await fn(*args, **kwargs)

    else:
        fn = cast(Callable[..., T], fn)

        @functools.wraps(fn)
        def wrapper(*args: object, **kwargs: object) -> T:
            with namespace_context(ns):
                return fn(*args, **kwargs)

    return wrapper