A module for Streamlit to set up traps, which are triggered on next script rerun (i.e. user action) if the trap wasn't reached.

st_traps.py

""" A module for Streamlit to set up traps, which are triggered on next script rerun (i.e. user action) if the trap wasn't reached.

Purpose:
    Help with session cleaning on dynamic forms. If user didn't reach it - then clear session.
"""
from collections.abc import Callable

import streamlit as st

__SESSION_STATE_DICT_KEY = "__INT_TRAPS_DICT"

def clear_session(prefix:str) -> None:
    """ A function to clear session state from values which keys start with the prefix.

    Args:
        prefix (str): a prefix to check for
    """
    for key in st.session_state:
        if key.startswith(prefix):
            st.session_state.pop(key)

def disarm() -> None:
    """ Disarm previously set traps.
    
    Ideally it should be one of last commands. If keys used in arming are unique enough - 
    there should be no overlapping with streamlit components' keys. But if you see any 
    issues - place it as one of the first commands. """
    traps:dict[str,tuple[int, Callable]] = st.session_state.get(__SESSION_STATE_DICT_KEY, {})
    # disarm those which have counter of 0
    to_disarm = {k for k,v in traps.items() if v[0]==0}
    for key in to_disarm:
        traps[key][1]()
        traps.pop(key, None)
    # rearm = reset counters to 0
    for key in traps:
        traps[key] = (0, traps[key][1])


def arm(key: str, disarm: Callable[[], None]) -> None:
    """ Arm a trap using a key (prefix), with a function to do the disarming 
    (like, clear session state of all keys with the prefix; actual function is up to you)

    Args:
        key (str): unique key for this section
        disarm (Callable[[], None]): a function to be called for cleaning (for example, clear_session(prefix=key))
    """
    assert key, "key cannot be empty"
    # disarm shouldn't be None here, but better to ensure that right here
    assert disarm and callable(disarm), "disarm must be callable" # type: ignore
    
    traps:dict[str,tuple[int, Callable]] = st.session_state.setdefault(__SESSION_STATE_DICT_KEY, {})
    traps[key] = (traps.get(key, (0, None))[0] + 1, disarm)