Update utils/storage.py

utils/storage.py

@@ -1,157 +1,300 @@
+"""
+Storage utilities for data persistence
+"""
 import json
 import pickle
-import pandas as pd
+import os
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
+import pandas as pd
 import streamlit as st
-
 from utils.error_handling import handle_data_exceptions, DataError, ValidationError
-from utils.logging import get_logger
 
-logger = get_logger(__name__)
 
-# In-memory storage for Streamlit (since file operations may be limited)
-_memory_store = {}
+# Default storage directories
+DATA_DIR = Path("data")
+CACHE_DIR = Path("cache")
+
+# Create directories if they don't exist
+DATA_DIR.mkdir(exist_ok=True)
+CACHE_DIR.mkdir(exist_ok=True)
 
-@handle_data_exceptions
-def save_data(key: str, data: Any, storage_type: str = "memory") -> bool:
-    """Save data with specified storage type"""
-    try:
-        if storage_type == "memory":
-            _memory_store[key] = data
-            logger.info(f"Data saved to memory with key: {key}")
-            return True
-        elif storage_type == "session":
-            if 'storage' not in st.session_state:
-                st.session_state.storage = {}
-            st.session_state.storage[key] = data
-            logger.info(f"Data saved to session state with key: {key}")
-            return True
-        else:
-            raise ValidationError(f"Unsupported storage type: {storage_type}")
-    except Exception as e:
-        logger.error(f"Failed to save data with key {key}: {str(e)}")
-        return False
 
 @handle_data_exceptions
-def load_data(key: str, storage_type: str = "memory", default_value: Any = None) -> Any:
-    """Load data from specified storage type"""
+def load_data(filename: str, data_dir: Optional[Path] = None) -> Optional[Dict[str, Any]]:
+    """
+    Load JSON data from file
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    filepath = data_dir / filename
+    
+    if not filepath.exists():
+        raise DataError(f"File {filepath} does not exist")
+    
     try:
-        if storage_type == "memory":
-            data = _memory_store.get(key, default_value)
-            if data is not None:
-                logger.info(f"Data loaded from memory with key: {key}")
-            return data
-        elif storage_type == "session":
-            if 'storage' not in st.session_state:
-                st.session_state.storage = {}
-            data = st.session_state.storage.get(key, default_value)
-            if data is not None:
-                logger.info(f"Data loaded from session state with key: {key}")
-            return data
-        else:
-            raise ValidationError(f"Unsupported storage type: {storage_type}")
+        with open(filepath, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        return data
+    except json.JSONDecodeError as e:
+        raise DataError(f"Invalid JSON in file {filepath}: {str(e)}")
     except Exception as e:
-        logger.error(f"Failed to load data with key {key}: {str(e)}")
-        return default_value
+        raise DataError(f"Error loading data from {filepath}: {str(e)}")
+
 
 @handle_data_exceptions
-def delete_data(key: str, storage_type: str = "memory") -> bool:
-    """Delete data from specified storage type"""
+def save_data(data: Dict[str, Any], filename: str, data_dir: Optional[Path] = None) -> bool:
+    """
+    Save data to JSON file
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    if not isinstance(data, dict):
+        raise ValidationError("Data must be a dictionary")
+    
+    filepath = data_dir / filename
+    
+    # Create directory if it doesn't exist
+    filepath.parent.mkdir(parents=True, exist_ok=True)
+    
     try:
-        if storage_type == "memory":
-            if key in _memory_store:
-                del _memory_store[key]
-                logger.info(f"Data deleted from memory with key: {key}")
-                return True
-        elif storage_type == "session":
-            if 'storage' in st.session_state and key in st.session_state.storage:
-                del st.session_state.storage[key]
-                logger.info(f"Data deleted from session state with key: {key}")
-                return True
-        return False
+        with open(filepath, 'w', encoding='utf-8') as f:
+            json.dump(data, f, indent=2, ensure_ascii=False)
+        return True
     except Exception as e:
-        logger.error(f"Failed to delete data with key {key}: {str(e)}")
-        return False
+        raise DataError(f"Error saving data to {filepath}: {str(e)}")
+
 
 @handle_data_exceptions
-def list_keys(storage_type: str = "memory") -> list:
-    """List all keys in specified storage type"""
+def load_dataframe(filename: str, data_dir: Optional[Path] = None) -> Optional[pd.DataFrame]:
+    """
+    Load DataFrame from various file formats
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    filepath = data_dir / filename
+    
+    if not filepath.exists():
+        raise DataError(f"File {filepath} does not exist")
+    
+    file_extension = filepath.suffix.lower()
+    
     try:
-        if storage_type == "memory":
-            return list(_memory_store.keys())
-        elif storage_type == "session":
-            if 'storage' not in st.session_state:
-                st.session_state.storage = {}
-            return list(st.session_state.storage.keys())
+        if file_extension == '.csv':
+            df = pd.read_csv(filepath)
+        elif file_extension == '.xlsx':
+            df = pd.read_excel(filepath)
+        elif file_extension == '.json':
+            df = pd.read_json(filepath)
+        elif file_extension == '.pkl':
+            df = pd.read_pickle(filepath)
         else:
-            raise ValidationError(f"Unsupported storage type: {storage_type}")
+            raise DataError(f"Unsupported file format: {file_extension}")
+        
+        return df
     except Exception as e:
-        logger.error(f"Failed to list keys for storage type {storage_type}: {str(e)}")
-        return []
+        raise DataError(f"Error loading dataframe from {filepath}: {str(e)}")
+
 
 @handle_data_exceptions
-def clear_storage(storage_type: str = "memory") -> bool:
-    """Clear all data from specified storage type"""
+def save_dataframe(df: pd.DataFrame, filename: str, data_dir: Optional[Path] = None, 
+                  file_format: str = 'csv') -> bool:
+    """
+    Save DataFrame to file
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    if not isinstance(df, pd.DataFrame):
+        raise ValidationError("Data must be a pandas DataFrame")
+    
+    if df.empty:
+        raise ValidationError("DataFrame cannot be empty")
+    
+    # Ensure filename has correct extension
+    if not filename.endswith(f'.{file_format}'):
+        filename = f"{filename}.{file_format}"
+    
+    filepath = data_dir / filename
+    
+    # Create directory if it doesn't exist
+    filepath.parent.mkdir(parents=True, exist_ok=True)
+    
     try:
-        if storage_type == "memory":
-            _memory_store.clear()
-            logger.info("Memory storage cleared")
-            return True
-        elif storage_type == "session":
-            if 'storage' in st.session_state:
-                st.session_state.storage.clear()
-            logger.info("Session storage cleared")
-            return True
+        if file_format == 'csv':
+            df.to_csv(filepath, index=False)
+        elif file_format == 'xlsx':
+            df.to_excel(filepath, index=False)
+        elif file_format == 'json':
+            df.to_json(filepath, orient='records', indent=2)
+        elif file_format == 'pkl':
+            df.to_pickle(filepath)
         else:
-            raise ValidationError(f"Unsupported storage type: {storage_type}")
+            raise DataError(f"Unsupported file format: {file_format}")
+        
+        return True
     except Exception as e:
-        logger.error(f"Failed to clear storage type {storage_type}: {str(e)}")
-        return False
+        raise DataError(f"Error saving dataframe to {filepath}: {str(e)}")
+
 
 @handle_data_exceptions
-def save_dataframe(df: pd.DataFrame, key: str, storage_type: str = "memory") -> bool:
-    """Save pandas DataFrame"""
-    if not isinstance(df, pd.DataFrame):
-        raise ValidationError("Data must be a pandas DataFrame")
+def load_pickle(filename: str, data_dir: Optional[Path] = None) -> Optional[Any]:
+    """
+    Load pickled data from file
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    filepath = data_dir / filename
     
-    # Convert DataFrame to dict for JSON serialization
-    df_dict = {
-        'data': df.to_dict('records'),
-        'columns': df.columns.tolist(),
-        'index': df.index.tolist()
-    }
+    if not filepath.exists():
+        raise DataError(f"File {filepath} does not exist")
     
-    return save_data(key, df_dict, storage_type)
+    try:
+        with open(filepath, 'rb') as f:
+            data = pickle.load(f)
+        return data
+    except Exception as e:
+        raise DataError(f"Error loading pickle from {filepath}: {str(e)}")
+
 
 @handle_data_exceptions
-def load_dataframe(key: str, storage_type: str = "memory") -> Optional[pd.DataFrame]:
-    """Load pandas DataFrame"""
-    df_dict = load_data(key, storage_type)
+def save_pickle(data: Any, filename: str, data_dir: Optional[Path] = None) -> bool:
+    """
+    Save data to pickle file
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
     
-    if df_dict is None:
-        return None
+    if not filename.endswith('.pkl'):
+        filename = f"{filename}.pkl"
+    
+    filepath = data_dir / filename
+    
+    # Create directory if it doesn't exist
+    filepath.parent.mkdir(parents=True, exist_ok=True)
     
     try:
-        df = pd.DataFrame(df_dict['data'])
-        if 'columns' in df_dict:
-            df.columns = df_dict['columns']
-        if 'index' in df_dict:
-            df.index = df_dict['index']
-        return df
+        with open(filepath, 'wb') as f:
+            pickle.dump(data, f)
+        return True
     except Exception as e:
-        logger.error(f"Failed to reconstruct DataFrame from stored data: {str(e)}")
-        return None
+        raise DataError(f"Error saving pickle to {filepath}: {str(e)}")
 
-def get_storage_info(storage_type: str = "memory") -> Dict[str, Any]:
-    """Get information about storage usage"""
+
+def list_files(data_dir: Optional[Path] = None, pattern: str = "*") -> list[Path]:
+    """
+    List files in data directory
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    if not data_dir.exists():
+        return []
+    
+    return list(data_dir.glob(pattern))
+
+
+def file_exists(filename: str, data_dir: Optional[Path] = None) -> bool:
+    """
+    Check if file exists
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    filepath = data_dir / filename
+    return filepath.exists()
+
+
+def delete_file(filename: str, data_dir: Optional[Path] = None) -> bool:
+    """
+    Delete file
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    filepath = data_dir / filename
+    
+    if filepath.exists():
+        try:
+            filepath.unlink()
+            return True
+        except Exception as e:
+            st.error(f"Error deleting file {filepath}: {str(e)}")
+            return False
+    return False
+
+
+def get_file_size(filename: str, data_dir: Optional[Path] = None) -> Optional[int]:
+    """
+    Get file size in bytes
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    filepath = data_dir / filename
+    
+    if filepath.exists():
+        return filepath.stat().st_size
+    return None
+
+
+def create_backup(filename: str, data_dir: Optional[Path] = None) -> bool:
+    """
+    Create backup of file
+    """
+    if data_dir is None:
+        data_dir = DATA_DIR
+    
+    filepath = data_dir / filename
+    
+    if not filepath.exists():
+        return False
+    
+    backup_filename = f"{filepath.stem}_backup{filepath.suffix}"
+    backup_filepath = data_dir / backup_filename
+    
     try:
-        keys = list_keys(storage_type)
-        return {
-            'storage_type': storage_type,
-            'total_keys': len(keys),
-            'keys': keys
-        }
+        import shutil
+        shutil.copy2(filepath, backup_filepath)
+        return True
     except Exception as e:
-        logger.error(f"Failed to get storage info: {str(e)}")
-        return {'storage_type': storage_type, 'total_keys': 0, 'keys': []}
+        st.error(f"Error creating backup: {str(e)}")
+        return False
+
+
+# Session state management
+def init_session_state(key: str, default_value: Any = None):
+    """
+    Initialize session state variable if it doesn't exist
+    """
+    if key not in st.session_state:
+        st.session_state[key] = default_value
+
+
+def get_session_state(key: str, default_value: Any = None) -> Any:
+    """
+    Get value from session state
+    """
+    return st.session_state.get(key, default_value)
+
+
+def set_session_state(key: str, value: Any):
+    """
+    Set value in session state
+    """
+    st.session_state[key] = value
+
+
+def clear_session_state(key: Optional[str] = None):
+    """
+    Clear session state (specific key or all)
+    """
+    if key:
+        if key in st.session_state:
+            del st.session_state[key]
+    else:
+        st.session_state.clear()