logger.py

import streamlit as st
from io import StringIO
import contextlib
import sys

# Redirect stdout and stderr
original_stdout = sys.stdout
original_stderr = sys.stderr
sys.stdout = StringIO()
sys.stderr = StringIO()

@st.cache
def capture_logs(func):
    with contextlib.redirect_stdout(sys.stdout), contextlib.redirect_stderr(sys.stderr):
        func()

# Get captured log messages
def get_captured_logs():
    captured_stdout = sys.stdout.getvalue()
    captured_stderr = sys.stderr.getvalue()
    return captured_stdout, captured_stderr

# Restore original stdout and stderr
def restore_original_streams():
    sys.stdout = original_stdout
    sys.stderr = original_stderr