huggan/sefa

1"""Adds pre-session state to StreamLit.2 3This file is borrowed from4https://gist.github.com/tvst/036da038ab3e999a64497f42de966a925"""6 7# pylint: disable=protected-access8 9try:10    import streamlit.ReportThread as ReportThread11    from streamlit.server.Server import Server12except ModuleNotFoundError:13    # Streamlit >= 0.65.014    import streamlit.report_thread as ReportThread15    from streamlit.server.server import Server16 17 18class SessionState(object):19    """Hack to add per-session state to Streamlit.20 21    Usage22    -----23 24    >>> import SessionState25    >>>26    >>> session_state = SessionState.get(user_name='', favorite_color='black')27    >>> session_state.user_name28    ''29    >>> session_state.user_name = 'Mary'30    >>> session_state.favorite_color31    'black'32 33    Since you set user_name above, next time your script runs this will be the34    result:35    >>> session_state = get(user_name='', favorite_color='black')36    >>> session_state.user_name37    'Mary'38 39    """40 41    def __init__(self, **kwargs):42        """A new SessionState object.43 44        Parameters45        ----------46        **kwargs : any47            Default values for the session state.48 49        Example50        -------51        >>> session_state = SessionState(user_name='', favorite_color='black')52        >>> session_state.user_name = 'Mary'53        ''54        >>> session_state.favorite_color55        'black'56 57        """58        for key, val in kwargs.items():59            setattr(self, key, val)60 61 62def get(**kwargs):63    """Gets a SessionState object for the current session.64 65    Creates a new object if necessary.66 67    Parameters68    ----------69    **kwargs : any70        Default values you want to add to the session state, if we're creating a71        new one.72 73    Example74    -------75    >>> session_state = get(user_name='', favorite_color='black')76    >>> session_state.user_name77    ''78    >>> session_state.user_name = 'Mary'79    >>> session_state.favorite_color80    'black'81 82    Since you set user_name above, next time your script runs this will be the83    result:84    >>> session_state = get(user_name='', favorite_color='black')85    >>> session_state.user_name86    'Mary'87 88    """89    # Hack to get the session object from Streamlit.90 91    ctx = ReportThread.get_report_ctx()92 93    this_session = None94 95    current_server = Server.get_current()96    if hasattr(current_server, '_session_infos'):97        # Streamlit < 0.5698        session_infos = Server.get_current()._session_infos.values()99    else:100        session_infos = Server.get_current()._session_info_by_id.values()101 102    for session_info in session_infos:103        s = session_info.session104        if (105            # Streamlit < 0.54.0106            (hasattr(s, '_main_dg') and s._main_dg == ctx.main_dg)107            or108            # Streamlit >= 0.54.0109            (not hasattr(s, '_main_dg') and s.enqueue == ctx.enqueue)110            or111            # Streamlit >= 0.65.2112            (not hasattr(s, '_main_dg') and113             s._uploaded_file_mgr == ctx.uploaded_file_mgr)114        ):115            this_session = s116 117    if this_session is None:118        raise RuntimeError(119            "Oh noes. Couldn't get your Streamlit Session object. "120            'Are you doing something fancy with threads?')121 122    # Got the session object! Now let's attach some state into it.123 124    if not hasattr(this_session, '_custom_session_state'):125        this_session._custom_session_state = SessionState(**kwargs)126 127    return this_session._custom_session_state128 129# pylint: enable=protected-access130