OpenBB/cli/openbb_cli/models/settings.py

"""Settings model."""

from enum import Enum
from typing import Any, Literal
from zoneinfo import available_timezones

from dotenv import dotenv_values, set_key
from openbb_core.app.version import get_package_version
from pydantic import BaseModel, ConfigDict, Field, model_validator

all_timezones = sorted(available_timezones())

from openbb_cli.config.constants import AVAILABLE_FLAIRS, ENV_FILE_SETTINGS

VERSION = get_package_version("openbb-cli")


class OutputMode(str, Enum):
    """Output mode for displaying results."""

    rich = "rich"
    json = "json"
    tsv = "tsv"
    html = "html"


class SettingGroups(Enum):
    """Setting types."""

    feature_flags = "feature_flag"
    preferences = "preference"


class Settings(BaseModel):
    """Settings model."""

    VERSION: str = VERSION

    TEST_MODE: bool = False
    DEBUG_MODE: bool = False
    DEV_BACKEND: bool = False

    PREVIOUS_USE: bool = False

    FILE_OVERWRITE: bool = Field(
        default=False,
        description="whether to overwrite Excel files if they already exists",
        json_schema_extra={
            "command": "overwrite",
            "group": SettingGroups.feature_flags.value,
        },
    )
    SHOW_VERSION: bool = Field(
        default=True,
        description="whether to show the version in the bottom right corner",
        json_schema_extra={
            "command": "version",
            "group": SettingGroups.feature_flags.value,
        },
    )
    USE_INTERACTIVE_DF: bool = Field(
        default=False,
        description="display tables in interactive window (when available)",
        json_schema_extra={
            "command": "interactive",
            "group": SettingGroups.feature_flags.value,
        },
    )
    OUTPUT_MODE: Literal["rich", "json", "tsv", "html"] = Field(
        default="tsv",
        description="output display mode (rich=terminal table, json=JSON, tsv=DataFrame string, html=browser)",
        json_schema_extra={
            "command": "output",
            "group": SettingGroups.preferences.value,
        },
    )
    USE_CLEAR_AFTER_CMD: bool = Field(
        default=False,
        description="clear console after each command",
        json_schema_extra={
            "command": "cls",
            "group": SettingGroups.feature_flags.value,
        },
    )
    USE_DATETIME: bool = Field(
        default=True,
        description="whether to show the date and time before the flair",
        json_schema_extra={
            "command": "datetime",
            "group": SettingGroups.feature_flags.value,
        },
    )
    USE_PROMPT_TOOLKIT: bool = Field(
        default=True,
        description="enable prompt toolkit (autocomplete and history)",
        json_schema_extra={
            "command": "promptkit",
            "group": SettingGroups.feature_flags.value,
        },
    )
    ENABLE_EXIT_AUTO_HELP: bool = Field(
        default=True,
        description="automatically print help when quitting menu",
        json_schema_extra={
            "command": "exithelp",
            "group": SettingGroups.feature_flags.value,
        },
    )
    ENABLE_RICH_PANEL: bool = Field(
        default=True,
        description="enable colorful rich CLI panel",
        json_schema_extra={
            "command": "richpanel",
            "group": SettingGroups.feature_flags.value,
        },
    )
    TOOLBAR_HINT: bool = Field(
        default=True,
        description="displays usage hints in the bottom toolbar",
        json_schema_extra={
            "command": "tbhint",
            "group": SettingGroups.feature_flags.value,
        },
    )
    SHOW_MSG_OBBJECT_REGISTRY: bool = Field(
        default=False,
        description="show obbject registry message after a new result is added",
        json_schema_extra={
            "command": "obbject_msg",
            "group": SettingGroups.feature_flags.value,
        },
    )

    TIMEZONE: Literal[tuple(all_timezones)] = Field(  # ty: ignore[invalid-type-form]
        default="America/New_York",
        description="pick timezone",
        json_schema_extra={
            "command": "timezone",
            "group": SettingGroups.preferences.value,
        },
    )
    FLAIR: Literal[tuple(AVAILABLE_FLAIRS)] = Field(  # ty: ignore[invalid-type-form]
        default=":openbb",
        description="choose flair icon",
        json_schema_extra={
            "command": "flair",
            "group": SettingGroups.preferences.value,
        },
    )
    N_TO_KEEP_OBBJECT_REGISTRY: int = Field(
        default=10,
        description="define the maximum number of obbjects allowed in the registry",
        json_schema_extra={
            "command": "obbject_res",
            "group": SettingGroups.preferences.value,
        },
    )
    N_TO_DISPLAY_OBBJECT_REGISTRY: int = Field(
        default=5,
        description="define the maximum number of cached results to display on the help menu",
        json_schema_extra={
            "command": "obbject_display",
            "group": SettingGroups.preferences.value,
        },
    )
    RICH_STYLE: str = Field(
        default="dark",
        description="apply a custom rich style to the CLI",
        json_schema_extra={
            "command": "console_style",
            "group": SettingGroups.preferences.value,
        },
    )
    ALLOWED_NUMBER_OF_ROWS: int = Field(
        default=20,
        description="number of rows to show in rich table mode (does not apply to json/stdio/html modes)",
        json_schema_extra={
            "command": "n_rows",
            "group": SettingGroups.preferences.value,
        },
    )
    ALLOWED_NUMBER_OF_COLUMNS: int = Field(
        default=5,
        description="number of columns to show in rich table mode (does not apply to json/stdio/html modes)",
        json_schema_extra={
            "command": "n_cols",
            "group": SettingGroups.preferences.value,
        },
    )

    model_config = ConfigDict(validate_assignment=True)

    def __repr__(self) -> str:
        """Return a string representation of the model."""
        return f"{self.__class__.__name__}\n\n" + "\n".join(
            f"{k}: {v}" for k, v in self.model_dump().items()
        )

    @model_validator(mode="before")
    @classmethod
    def from_env(cls, values: dict) -> dict:
        """Load settings from .env."""
        settings = {}
        settings.update(dotenv_values(ENV_FILE_SETTINGS))
        settings.update(values)
        filtered = {k.replace("OPENBB_", ""): v for k, v in settings.items()}
        return filtered

    def set_item(self, key: str, value: Any) -> None:
        """Set an item in the model and save to .env."""
        setattr(self, key, value)
        set_key(str(ENV_FILE_SETTINGS), "OPENBB_" + key, str(value))