File size: 1,785 Bytes
51545af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """
Progress management utilities for Streamlit UI.
"""
from abc import ABC, abstractmethod
from typing import Optional, Callable
import streamlit as st
class ProgressContext(ABC):
"""Base class for different progress UI patterns"""
@abstractmethod
def __enter__(self) -> Callable[[float, str], None]:
pass
@abstractmethod
def __exit__(self, exc_type, exc_val, exc_tb):
pass
class StreamlitProgressContext(ProgressContext):
"""Standard Streamlit progress bar with automatic cleanup"""
def __init__(self, placeholder, success_message: Optional[str] = None):
self.placeholder = placeholder
self.success_message = success_message
self.progress_bar = None
def __enter__(self):
self.progress_bar = self.placeholder.progress(0, text="Starting...")
return self.update_progress
def update_progress(self, progress: float, text: str):
if self.progress_bar:
self.progress_bar.progress(progress, text=text)
def __exit__(self, exc_type, exc_val, exc_tb):
if self.progress_bar:
self.progress_bar.empty()
if exc_type is None and self.success_message:
self.placeholder.success(self.success_message)
elif exc_type is not None:
self.placeholder.error(f"Error: {exc_val}")
class MockProgressContext(ProgressContext):
"""Mock progress context for testing - captures progress updates without UI"""
def __init__(self):
self.updates = []
def __enter__(self):
return self.capture_progress
def capture_progress(self, progress: float, text: str):
self.updates.append((progress, text))
def __exit__(self, *args):
pass
|