Update utils/storage.py

utils/storage.py

@@ -1,239 +1,171 @@
 import json
 import os
+import pickle
 from pathlib import Path
-from typing import Dict, List, Any, Union, Optional
+from typing import Any, Dict, Optional, Union
+import pandas as pd
 
-from utils.config import FILE_PATHS, DATA_DIR, EXPORT_DIR
-from utils.logging import get_logger, log_performance
 from utils.error_handling import handle_data_exceptions, DataError, ValidationError
+from utils.logging import get_logger
 
-# Initialize logger
 logger = get_logger(__name__)
 
-@handle_data_exceptions
-def load_data(file_path: Union[str, Path], default: Any = None) -> Any:
-    """
-    Load data from a JSON file
-    
-    Args:
-        file_path: Path to the JSON file
-        default: Default value to return if file doesn't exist
-        
-    Returns:
-        Data from the JSON file or default value
-        
-    Raises:
-        DataError: If there's an error loading the data
-    """
-    file_path = Path(file_path)
-    
-    if not file_path.exists():
-        logger.debug(f"File not found: {file_path}. Returning default value.")
-        return default if default is not None else {}
-    
-    try:
-        logger.debug(f"Loading data from {file_path}")
-        with open(file_path, 'r', encoding='utf-8') as f:
-            data = json.load(f)
-        logger.debug(f"Successfully loaded data from {file_path}")
-        return data
-    except json.JSONDecodeError as e:
-        logger.error(f"JSON decode error in {file_path}: {str(e)}")
-        raise DataError(f"Invalid JSON format in {file_path}", {"original_error": str(e)}) from e
-    except FileNotFoundError as e:
-        logger.error(f"File not found: {file_path}")
-        return default if default is not None else {}
-    except Exception as e:
-        logger.error(f"Unexpected error loading {file_path}: {str(e)}")
-        raise DataError(f"Error loading data from {file_path}", {"original_error": str(e)}) from e
+# Create data directory if it doesn't exist
+DATA_DIR = Path("data")
+DATA_DIR.mkdir(exist_ok=True)
 
 @handle_data_exceptions
-def save_data(file_path: Union[str, Path], data: Any) -> bool:
+def save_data(data: Any, filename: str, format_type: str = "json") -> bool:
     """
-    Save data to a JSON file
+    Save data to file
     
     Args:
-        file_path: Path to the JSON file
         data: Data to save
-        
+        filename: Name of the file
+        format_type: Format to save ('json', 'pickle', 'csv', 'excel')
+    
     Returns:
-        True if successful, False otherwise
-        
-    Raises:
-        DataError: If there's an error saving the data
+        bool: True if successful, False otherwise
     """
-    file_path = Path(file_path)
-    
-    # Create directory if it doesn't exist
-    file_path.parent.mkdir(parents=True, exist_ok=True)
-    
     try:
-        logger.debug(f"Saving data to {file_path}")
-        with open(file_path, 'w', encoding='utf-8') as f:
-            json.dump(data, f, indent=2, ensure_ascii=False)
-        logger.debug(f"Successfully saved data to {file_path}")
+        filepath = DATA_DIR / filename
+        
+        if format_type == "json":
+            with open(filepath, 'w', encoding='utf-8') as f:
+                json.dump(data, f, indent=2, ensure_ascii=False, default=str)
+        
+        elif format_type == "pickle":
+            with open(filepath, 'wb') as f:
+                pickle.dump(data, f)
+        
+        elif format_type == "csv":
+            if isinstance(data, pd.DataFrame):
+                data.to_csv(filepath, index=False)
+            else:
+                raise DataError("CSV format requires pandas DataFrame")
+        
+        elif format_type == "excel":
+            if isinstance(data, pd.DataFrame):
+                data.to_excel(filepath, index=False)
+            else:
+                raise DataError("Excel format requires pandas DataFrame")
+        
+        else:
+            raise DataError(f"Unsupported format: {format_type}")
+        
+        logger.info(f"Data saved successfully to {filepath}")
         return True
+        
     except Exception as e:
-        logger.error(f"Error saving data to {file_path}: {str(e)}")
-        raise DataError(f"Error saving data to {file_path}", {"original_error": str(e)}) from e
+        logger.error(f"Failed to save data to {filename}: {str(e)}")
+        raise DataError(f"Failed to save data: {str(e)}")
 
 @handle_data_exceptions
-def export_data(data: Any, format_type: str = 'json', file_path: Optional[Union[str, Path]] = None) -> Union[str, bool]:
+def load_data(filename: str, format_type: str = "json") -> Any:
     """
-    Export data to various formats
+    Load data from file
     
     Args:
-        data: Data to export
-        format_type: Format to export to ('json', 'csv', 'markdown')
-        file_path: Path to save the exported data (optional)
-        
+        filename: Name of the file
+        format_type: Format to load ('json', 'pickle', 'csv', 'excel')
+    
     Returns:
-        Exported data as string if file_path is None, otherwise True if successful
-        
-    Raises:
-        ValidationError: If the data format is invalid
-        DataError: If there's an error exporting the data
+        Any: Loaded data or None if failed
     """
-    logger.debug(f"Exporting data to {format_type} format")
-    
-    # Validate format type
-    if format_type not in ['json', 'csv', 'markdown']:
-        logger.error(f"Unsupported format type: {format_type}")
-        raise ValidationError(f"Unsupported format type: {format_type}")
-    
     try:
-        if format_type == 'json':
-            result = json.dumps(data, indent=2, ensure_ascii=False)
-        elif format_type == 'csv':
-            # Simple CSV conversion for list of dictionaries
-            if not isinstance(data, list) or not all(isinstance(item, dict) for item in data):
-                logger.error("CSV export only supports list of dictionaries")
-                raise ValidationError("CSV export only supports list of dictionaries")
-            
-            if not data:
-                return ""
-                
-            headers = list(data[0].keys())
-            result = ",".join(headers) + "\n"
-            
-            for item in data:
-                row = ",".join([str(item.get(header, "")).replace(",", ";") for header in headers])
-                result += row + "\n"
-        elif format_type == 'markdown':
-            # Simple Markdown conversion for list of dictionaries
-            if not isinstance(data, list) or not all(isinstance(item, dict) for item in data):
-                logger.error("Markdown export only supports list of dictionaries")
-                raise ValidationError("Markdown export only supports list of dictionaries")
-                
-            if not data:
-                return ""
-                
-            headers = list(data[0].keys())
-            result = "| " + " | ".join(headers) + " |\n"
-            result += "| " + " | ".join(["---" for _ in headers]) + " |\n"
-            
-            for item in data:
-                row = "| " + " | ".join([str(item.get(header, "")).replace("|", "\|") for header in headers]) + " |"
-                result += row + "\n"
-    except Exception as e:
-        logger.error(f"Error formatting data for export: {str(e)}")
-        raise DataError(f"Error formatting data for export", {"original_error": str(e)}) from e
-    
-    if file_path is None:
-        return result
-    
-    file_path = Path(file_path)
-    file_path.parent.mkdir(parents=True, exist_ok=True)
-    
-    try:
-        logger.debug(f"Writing exported data to {file_path}")
-        with open(file_path, 'w', encoding='utf-8') as f:
-            f.write(result)
-        logger.debug(f"Successfully exported data to {file_path}")
-        return True
+        filepath = DATA_DIR / filename
+        
+        if not filepath.exists():
+            logger.warning(f"File not found: {filepath}")
+            return None
+        
+        if format_type == "json":
+            with open(filepath, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+        
+        elif format_type == "pickle":
+            with open(filepath, 'rb') as f:
+                data = pickle.load(f)
+        
+        elif format_type == "csv":
+            data = pd.read_csv(filepath)
+        
+        elif format_type == "excel":
+            data = pd.read_excel(filepath)
+        
+        else:
+            raise DataError(f"Unsupported format: {format_type}")
+        
+        logger.info(f"Data loaded successfully from {filepath}")
+        return data
+        
     except Exception as e:
-        logger.error(f"Error writing exported data to {file_path}: {str(e)}")
-        raise DataError(f"Error exporting data to {file_path}", {"original_error": str(e)}) from e
+        logger.error(f"Failed to load data from {filename}: {str(e)}")
+        raise DataError(f"Failed to load data: {str(e)}")
 
 @handle_data_exceptions
-def get_file_path(data_type: str) -> Path:
+def delete_data(filename: str) -> bool:
     """
-    Get the file path for a specific data type
+    Delete data file
     
     Args:
-        data_type: Type of data (e.g., 'tasks', 'notes', 'goals')
-        
+        filename: Name of the file to delete
+    
     Returns:
-        Path object for the specified data type
-        
-    Raises:
-        ValidationError: If the data type is invalid
+        bool: True if successful, False otherwise
     """
-    if data_type not in FILE_PATHS:
-        logger.error(f"Invalid data type: {data_type}")
-        raise ValidationError(f"Invalid data type: {data_type}")
-    
-    return FILE_PATHS[data_type]
+    try:
+        filepath = DATA_DIR / filename
+        
+        if filepath.exists():
+            filepath.unlink()
+            logger.info(f"File deleted successfully: {filepath}")
+            return True
+        else:
+            logger.warning(f"File not found for deletion: {filepath}")
+            return False
+            
+    except Exception as e:
+        logger.error(f"Failed to delete file {filename}: {str(e)}")
+        raise DataError(f"Failed to delete file: {str(e)}")
 
 @handle_data_exceptions
-def create_backup(backup_name: str = None) -> Optional[Path]:
+def list_data_files() -> list:
     """
-    Create a backup of all data files
+    List all data files in the data directory
     
-    Args:
-        backup_name: Optional name for the backup file
-        
     Returns:
-        Path to the backup file if successful, None otherwise
+        list: List of filenames
     """
-    import zipfile
-    from datetime import datetime
-    
-    if backup_name is None:
-        backup_name = f"mona_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-    
-    backup_file = EXPORT_DIR / f"{backup_name}.zip"
-    
     try:
-        logger.info(f"Creating backup: {backup_file}")
-        with zipfile.ZipFile(backup_file, 'w', zipfile.ZIP_DEFLATED) as zipf:
-            for data_type, file_path in FILE_PATHS.items():
-                if file_path.exists():
-                    zipf.write(file_path, arcname=file_path.name)
-                    logger.debug(f"Added {file_path} to backup")
-        
-        logger.info(f"Backup created successfully: {backup_file}")
-        return backup_file
+        files = [f.name for f in DATA_DIR.iterdir() if f.is_file()]
+        logger.info(f"Found {len(files)} data files")
+        return files
+        
     except Exception as e:
-        logger.error(f"Error creating backup: {str(e)}")
-        raise DataError(f"Error creating backup", {"original_error": str(e)}) from e
+        logger.error(f"Failed to list data files: {str(e)}")
+        return []
 
-@handle_data_exceptions
-def restore_backup(backup_file: Union[str, Path]) -> bool:
+def get_file_size(filename: str) -> Optional[int]:
     """
-    Restore data from a backup file
+    Get file size in bytes
     
     Args:
-        backup_file: Path to the backup file
-        
+        filename: Name of the file
+    
     Returns:
-        True if successful, False otherwise
+        int: File size in bytes or None if file doesn't exist
     """
-    import zipfile
-    
-    backup_file = Path(backup_file)
-    
-    if not backup_file.exists():
-        logger.error(f"Backup file not found: {backup_file}")
-        raise DataError(f"Backup file not found: {backup_file}")
-    
     try:
-        logger.info(f"Restoring from backup: {backup_file}")
-        with zipfile.ZipFile(backup_file, 'r') as zipf:
-            zipf.extractall(DATA_DIR)
-        
-        logger.info(f"Backup restored successfully from {backup_file}")
-        return True
+        filepath = DATA_DIR / filename
+        if filepath.exists():
+            return filepath.stat().st_size
+        return None
     except Exception as e:
-        logger.error(f"Error restoring backup: {str(e)}")
-        raise DataError(f"Error restoring backup", {"original_error": str(e)}) from e
+        logger.error(f"Failed to get file size for {filename}: {str(e)}")
+        return None
+
+def ensure_data_directory():
+    """Ensure data directory exists"""
+    DATA_DIR.mkdir(exist_ok=True)
+    return DATA_DIR