diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index cfd580e89..5cbb41d06 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -15,7 +15,7 @@ from . import util from . import bolt_util from .plotter import Plotter -from .util import in_databricks, setup_logger, in_ipython, make_iframe +from .util import in_databricks, setup_logger, in_ipython, make_iframe, display_message_html from .exceptions import SsoRetrieveTokenTimeoutException, TokenExpireException from .messages import ( @@ -2494,7 +2494,10 @@ def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent if not PyGraphistry.sso_repeat_get_token(repeat, wait): msg_text = f'{msg_text}\nFailed to get token after {repeat*wait} seconds ....' if not fail_silent: - raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the login process") + msg = f"Failed to get token after {repeat*wait} seconds. Please re-run the login process" + if in_ipython() or in_databricks or PyGraphistry.set_sso_opt_into_type == "display": + display_message_html("{msg}") + raise Exception(msg) else: msg_text = f'{msg_text}\nGot token' print(msg_text) diff --git a/graphistry/util.py b/graphistry/util.py index e4b2b0a3c..ba47f55eb 100644 --- a/graphistry/util.py +++ b/graphistry/util.py @@ -172,6 +172,14 @@ def check_set_memoize( weakref[hashed] = w return False +def display_message_html(message: str, cleared: Optional[bool] = False): + from IPython.display import display, HTML, clear_output + + if cleared: + clear_output() + + display(HTML(message)) + def make_iframe_srcdoc(srcdoc: str): srcdoc = 'srcdoc="{srcdoc}" onload="this.removeAttribute(\'srcdoc\')"' return srcdoc