diff --git a/openbb_terminal/core/models/profile_model.py b/openbb_terminal/core/models/profile_model.py index 4fb75f1b588..9cbe215ddd2 100644 --- a/openbb_terminal/core/models/profile_model.py +++ b/openbb_terminal/core/models/profile_model.py @@ -14,8 +14,9 @@ class ProfileModel(BaseModel): uuid: str = "" email: str = "" username: str = "" + remember: bool = True - def load_user_info(self, session: dict, email: str): + def load_user_info(self, session: dict, email: str, remember: bool): """Load user info from login info. Parameters @@ -24,12 +25,15 @@ class ProfileModel(BaseModel): The login info. email : str The email. + remember : bool + Remember the session. """ self.token_type = session.get("token_type", "") self.token = session.get("access_token", "") self.uuid = session.get("uuid", "") self.email = email self.username = self.email[: self.email.find("@")] + self.remember = remember def get_uuid(self) -> str: """Get uuid. diff --git a/openbb_terminal/core/session/session_controller.py b/openbb_terminal/core/session/session_controller.py index d14021e610e..a1c18d1bb70 100644 --- a/openbb_terminal/core/session/session_controller.py +++ b/openbb_terminal/core/session/session_controller.py @@ -50,9 +50,11 @@ def get_user_input() -> Tuple[str, str, bool]: message="> Password: ", is_password=True, ) - save = s.prompt(message="> Remember me? (y/n): ", is_password=False).lower() == "y" + remember = ( + s.prompt(message="> Remember me? (y/n): ", is_password=False).lower() == "y" + ) - return email, password, save + return email, password, remember def prompt(welcome=True): @@ -69,12 +71,12 @@ def prompt(welcome=True): display_welcome_message() while True: - email, password, save = get_user_input() + email, password, remember = get_user_input() if not email: return launch_terminal() - session = create_session(email, password, save) + session = create_session(email, password, remember) if isinstance(session, dict) and session: - return login_and_launch(session=session) + return login_and_launch(session, remember) def launch_terminal(): @@ -85,15 +87,17 @@ def launch_terminal(): terminal_controller.parse_args_and_run() -def login_and_launch(session: dict): +def login_and_launch(session: dict, remember: bool = False): """Login and launch terminal. Parameters ---------- session : dict The session info. + remember : bool, optional + Remember the session, by default False """ - status = login(session) + status = login(session, remember) if status in [LoginStatus.SUCCESS, LoginStatus.NO_RESPONSE]: launch_terminal() elif status == LoginStatus.FAILED: @@ -108,7 +112,7 @@ def main(): if not local_session: prompt() else: - login_and_launch(session=local_session) + login_and_launch(session=local_session, remember=True) if __name__ == "__main__": diff --git a/openbb_terminal/core/session/session_model.py b/openbb_terminal/core/session/session_model.py index b6962707f48..f47acc69514 100644 --- a/openbb_terminal/core/session/session_model.py +++ b/openbb_terminal/core/session/session_model.py @@ -28,7 +28,7 @@ from openbb_terminal.core.session.sources_handler import get_updated_hub_sources from openbb_terminal.core.session.utils import run_thread from openbb_terminal.helper_funcs import system_clear from openbb_terminal.loggers import setup_logging -from openbb_terminal.rich_config import console +from openbb_terminal.rich_config import console, optional_rich_track # pylint: disable=consider-using-f-string @@ -76,13 +76,15 @@ def create_session_from_token(token: str, save: bool) -> Dict[Any, Any]: return session -def login(session: dict) -> LoginStatus: +def login(session: Dict, remember: bool = False) -> LoginStatus: """Login and load user info. Parameters ---------- session : dict The session info. + remember : bool, optional + Remember the session, by default False """ # Create a new user: # credentials: stored in hub, but we fallback to local (.env) @@ -101,11 +103,11 @@ def login(session: dict) -> LoginStatus: if response.status_code == 200: configs = json.loads(response.content) email = configs.get("email", "") - hub_user.profile.load_user_info(session, email) + hub_user.profile.load_user_info(session, email, remember) set_current_user(hub_user) auth_header = hub_user.profile.get_auth_header() - run_thread(download_and_save_routines, {"auth_header": auth_header}) + download_and_save_routines(auth_header) run_thread( update_backend_sources, {"auth_header": auth_header, "configs": configs} ) @@ -119,21 +121,23 @@ def login(session: dict) -> LoginStatus: return LoginStatus.NO_RESPONSE -def download_and_save_routines(auth_header: str, silent: bool = True): +def download_and_save_routines(auth_header: str): """Download and save routines. Parameters ---------- auth_header : str The authorization header, e.g. "Bearer ". - silent : bool - Whether to silence the console output, by default True """ - routines = download_routines(auth_header=auth_header, silent=silent) - for name, content in routines.items(): - save_routine( - file_name=f"{name}.openbb", routine=content, force=True, silent=silent - ) + try: + console.print("Downloading routines...") + routines = download_routines(auth_header=auth_header) + for name, content in optional_rich_track( + routines.items(), desc="Saving routines" + ): + save_routine(file_name=f"{name}.openbb", routine=content, force=True) + except Exception: + console.print("[red]Failed to download and save routines.[/red]") def update_backend_sources(auth_header, configs, silent: bool = True): diff --git a/openbb_terminal/parent_classes.py b/openbb_terminal/parent_classes.py index 26c833791dc..39095249787 100644 --- a/openbb_terminal/parent_classes.py +++ b/openbb_terminal/parent_classes.py @@ -21,8 +21,11 @@ from prompt_toolkit.formatted_text import HTML from prompt_toolkit.styles import Style from rich.markdown import Markdown +import openbb_terminal.core.session.local_model as Local + # IMPORTS INTERNAL from openbb_terminal.core.completer.choices import build_controller_choice_map +from openbb_terminal.core.config.paths import HIST_FILE_PATH from openbb_terminal.core.session.current_user import get_current_user, is_local from openbb_terminal.cryptocurrency import cryptocurrency_helpers from openbb_terminal.custom_prompt_toolkit import NestedCompleter @@ -473,6 +476,11 @@ class BaseController(metaclass=ABCMeta): for _ in range(self.PATH.count("/")): self.queue.insert(0, "quit") + if not is_local(): + Local.remove(get_current_user().preferences.USER_ROUTINES_DIRECTORY / "hub") + if not get_current_user().profile.remember: + Local.remove(HIST_FILE_PATH) + @log_start_end(log=logger) def call_reset(self, _) -> None: """Process reset command. diff --git a/tests/openbb_terminal/session/test_sdk_session.py b/tests/openbb_terminal/session/test_sdk_session.py index 4c4066610ad..d5ebb7c9fe9 100644 --- a/tests/openbb_terminal/session/test_sdk_session.py +++ b/tests/openbb_terminal/session/test_sdk_session.py @@ -225,7 +225,7 @@ def test_login_no_args( def test_logout(mocker, test_user): - test_user.profile.load_user_info(TEST_SESSION, "testy@email.com") + test_user.profile.load_user_info(TEST_SESSION, "testy@email.com", False) mocker.patch( target="openbb_terminal.core.session.sdk_session.get_current_user", return_value=test_user, @@ -253,6 +253,7 @@ def test_whoami(mocker, test_user): "uuid": "MOCK_UUID", }, email="MOCK_EMAIL", + remember=False, ) mocker.patch( target="openbb_terminal.core.session.sdk_session.get_current_user", diff --git a/tests/openbb_terminal/session/test_user.py b/tests/openbb_terminal/session/test_user.py index 78ee5847077..656d17c2838 100644 --- a/tests/openbb_terminal/session/test_user.py +++ b/tests/openbb_terminal/session/test_user.py @@ -29,7 +29,7 @@ def fixture_test_user(): profile=ProfileModel(), sources=SourcesModel(), ) - test_user.profile.load_user_info(TEST_SESSION, "test@email.com") + test_user.profile.load_user_info(TEST_SESSION, "test@email.com", False) return test_user @@ -39,6 +39,7 @@ def test_load_user_info(test_user): assert test_user.profile.token_type == "bearer" assert test_user.profile.uuid == "test_uuid" assert test_user.profile.email == "test@email.com" + assert test_user.profile.remember is False def test_get_session(test_user):