An alternative API for st.session_state using a class decorator.

state_class_patch.py

"""An alternative API for st.session_state using a class decorator.

So instead of doing this...

    if "run_counter" in st.session_state:
      st.session_state.run_counter = 0

    st.session_state.run_counter += 1

    st.write("Number of reruns:", st.session_state.run_counter)
    st.button("Rerun")
 
 
...you do this:
    import state_class_patch 

    @st.session_state
    class UserSession:
      run_counter = 0

    UserSession.run_counter += 1

    st.write("Number of reruns:", UserSession.run_counter)
    st.button("Rerun")


This may not look like much of an improvement, but it becomes more useful
when you have multiple properties to instantiate...

    @st.session_state
    class UserInfo:
      name = "Alice"
      job = "Student"
      age = 33
    

...or when you want to more cleanly operate on your session variables via
methods:

    @st.session_state
    class Session:
      counter = 0

      def increment(self):
        self.counter += 1

      def reset(self):
        self.counter = 0

    st.metric("Counter", UserSession.counter)
    st.button("Increment", on_click=UserSession.increment)
    st.button("Reset", on_click=UserSession.reset)
"""
import streamlit as st
from dataclasses import dataclass

def session_state__call__(self, cls):
  """Wrapper around the dataclass decorator, which also adds the
  instance to st.session_state.
  """
  name = cls.__name__
  wrapped_cls = dataclass(cls)

  if name not in st.session_state:
    st.session_state[name] = wrapped_cls()

  return st.session_state[name]

# Monkey-patch st.session_state so it can be called like a function.
_SessionStateClass = type(st.session_state)
_SessionStateClass.__call__ = session_state__call__.__get__(st.session_state, _SessionStateClass)