blatted/blatted/tools/context.py

from contextvars import ContextVar
from enum import Enum

from rich.console import Console


class BlattedEnvironment(Enum):
    CLI = "console line interface"
    TUI = "terminal user interface"


blatted_environment_var: ContextVar = ContextVar("blatted_environment")

blatted_console_var: ContextVar = ContextVar("blatted_console")
try:
    # test if the variable was set
    blatted_console_var.get()
except LookupError:
    console = Console()
    blatted_console_var.set(console)


def set_environment(mode: BlattedEnvironment = BlattedEnvironment.CLI) -> None:
    blatted_environment_var.set(mode)


def get_environment() -> BlattedEnvironment:
    try:
        return blatted_environment_var.get()
    except LookupError:
        blatted_environment_var.set(BlattedEnvironment.CLI)
        return blatted_environment_var.get()


def get_console() -> Console:
    return blatted_console_var.get()