diff --git a/openbb_terminal/session/session_controller.py b/openbb_terminal/session/session_controller.py index cbab08d4655..288c4b78289 100644 --- a/openbb_terminal/session/session_controller.py +++ b/openbb_terminal/session/session_controller.py @@ -1,12 +1,13 @@ from typing import Tuple -import json +from openbb_terminal import terminal_controller import openbb_terminal.session.local_model as Local -import openbb_terminal.session.hub_model as Hub -from openbb_terminal.session.session_model import create_session -from openbb_terminal.session.user import User +from openbb_terminal.session.session_model import ( + LoginStatus, + create_session, + login, +) from openbb_terminal.core.config.paths import PACKAGE_DIRECTORY from openbb_terminal.rich_config import console -from openbb_terminal import terminal_controller def display_welcome_message(): @@ -34,8 +35,8 @@ def get_user_input() -> Tuple[str, str, bool]: return email, password, save -def login_prompt(welcome=True, guest_allowed=False): - """Login prompt and launch terminal if login is successful. +def prompt(welcome=True, guest_allowed=False): + """Prompt and launch terminal if login is successful. Parameters ---------- @@ -49,18 +50,20 @@ def login_prompt(welcome=True, guest_allowed=False): email, password, save = get_user_input() if not email or not password: if guest_allowed: - return terminal_controller.parse_args_and_run() + return launch_terminal() continue session = create_session(email, password, save) if isinstance(session, dict) and session: - break - - login(session=session) + return login_and_launch(session=session, guest_allowed=guest_allowed) -# TODO: Move some login inside this function to the session_model.py -def login(session: dict, guest_allowed: bool = True): +def launch_terminal(): + """Launch terminal""" + terminal_controller.parse_args_and_run() + + +def login_and_launch(session: dict, guest_allowed: bool = True): """Login and launch terminal. Parameters @@ -70,25 +73,28 @@ def login(session: dict, guest_allowed: bool = True): guest_allowed : bool, optional Allow guest login, by default True """ - response = Hub.fetch_user_configs(session) - if response: - if response.status_code == 200: - User.load_user_info(session) - Local.apply_configs(configs=json.loads(response.content)) - terminal_controller.parse_args_and_run() - else: - login_prompt(welcome=False, guest_allowed=guest_allowed) + status = login(session) + if status == LoginStatus.SUCCESS: + launch_terminal() + elif status == LoginStatus.FAILED: + prompt(welcome=False, guest_allowed=guest_allowed) else: - login_prompt(welcome=True, guest_allowed=guest_allowed) + prompt(welcome=True, guest_allowed=guest_allowed) def main(guest_allowed: bool = True): - """Main function""" + """Main function + + Parameters + ---------- + guest_allowed : bool, optional + Allow guest login, by default True + """ local_session = Local.get_session() if not local_session: - login_prompt(guest_allowed=guest_allowed) + prompt(guest_allowed=guest_allowed) else: - login(session=local_session, guest_allowed=guest_allowed) + login_and_launch(session=local_session, guest_allowed=guest_allowed) if __name__ == "__main__": diff --git a/openbb_terminal/session/session_model.py b/openbb_terminal/session/session_model.py index 4ff4244cd26..c639c820a2f 100644 --- a/openbb_terminal/session/session_model.py +++ b/openbb_terminal/session/session_model.py @@ -1,7 +1,9 @@ import importlib import os import sys +import json import matplotlib.pyplot as plt +from enum import Enum import openbb_terminal.session.local_model as Local import openbb_terminal.session.hub_model as Hub from openbb_terminal.session.user import User @@ -9,6 +11,12 @@ from openbb_terminal.helper_funcs import system_clear from openbb_terminal.rich_config import console +class LoginStatus(Enum): + SUCCESS = 1 + FAILED = 2 + NO_RESPONSE = 3 + + def create_session(email: str, password: str, save: bool) -> dict: """Create a session. @@ -28,6 +36,24 @@ def create_session(email: str, password: str, save: bool) -> dict: return session +def login(session: dict) -> LoginStatus: + """Login and load user info. + + Parameters + ---------- + session : dict + The session info. + """ + response = Hub.fetch_user_configs(session) + if response: + if response.status_code == 200: + User.load_user_info(session) + Local.apply_configs(configs=json.loads(response.content)) + return LoginStatus.SUCCESS + return LoginStatus.FAILED + return LoginStatus.NO_RESPONSE + + def logout(cls: bool = False): """Logout and clear session.