Update utils/storage.py

utils/storage.py

@@ -1,239 +1,399 @@
+"""
+Storage utilities for the application.
+Handles data persistence, loading, and saving operations.
+"""
+
 import json
+import pickle
 import os
+import shutil
 from pathlib import Path
-from typing import Dict, List, Any, Union, Optional
+from typing import Any, Dict, List, Optional, Union
+from datetime import datetime
+import tempfile
 
-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.error_handling import handle_data_exceptions, DataError, ValidationError, StorageError
 
-# Initialize logger
 logger = get_logger(__name__)
 
-@handle_data_exceptions
-def load_data(file_path: Union[str, Path], default: Any = None) -> Any:
+# Default storage directory
+DEFAULT_STORAGE_DIR = Path("data")
+DEFAULT_STORAGE_DIR.mkdir(exist_ok=True)
+
+@handle_data_exceptions(default_return=None)
+@log_performance
+def load_data(
+    filename: str,
+    storage_dir: Union[str, Path] = None,
+    format: str = "auto",
+    encoding: str = "utf-8",
+    default: Any = None
+) -> Any:
     """
-    Load data from a JSON file
+    Load data from a file.
     
     Args:
-        file_path: Path to the JSON file
+        filename: Name of the file to load
+        storage_dir: Directory containing the file (defaults to DEFAULT_STORAGE_DIR)
+        format: File format ('json', 'pickle', 'txt', or 'auto' to detect from extension)
+        encoding: Text encoding for text files
         default: Default value to return if file doesn't exist
-        
+    
     Returns:
-        Data from the JSON file or default value
-        
+        Loaded data or default value
+    
     Raises:
-        DataError: If there's an error loading the data
+        StorageError: If file cannot be loaded
     """
-    file_path = Path(file_path)
+    if storage_dir is None:
+        storage_dir = DEFAULT_STORAGE_DIR
+    
+    file_path = Path(storage_dir) / filename
     
+    # Return default if file doesn't exist
     if not file_path.exists():
-        logger.debug(f"File not found: {file_path}. Returning default value.")
-        return default if default is not None else {}
+        logger.info(f"File {file_path} does not exist, returning default value")
+        return default
+    
+    # Auto-detect format from extension
+    if format == "auto":
+        ext = file_path.suffix.lower()
+        if ext == ".json":
+            format = "json"
+        elif ext in [".pkl", ".pickle"]:
+            format = "pickle"
+        else:
+            format = "txt"
     
     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}")
+        logger.debug(f"Loading data from {file_path} with format {format}")
+        
+        if format == "json":
+            with open(file_path, 'r', encoding=encoding) as f:
+                data = json.load(f)
+        elif format == "pickle":
+            with open(file_path, 'rb') as f:
+                data = pickle.load(f)
+        elif format == "txt":
+            with open(file_path, 'r', encoding=encoding) as f:
+                data = f.read()
+        else:
+            raise StorageError(f"Unsupported format: {format}", operation="load")
+        
+        logger.info(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 {}
+        raise StorageError(f"Invalid JSON in file {file_path}: {e}", operation="load")
+    except pickle.UnpicklingError as e:
+        raise StorageError(f"Invalid pickle data in file {file_path}: {e}", operation="load")
+    except UnicodeDecodeError as e:
+        raise StorageError(f"Encoding error reading file {file_path}: {e}", operation="load")
     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
+        raise StorageError(f"Error loading file {file_path}: {e}", operation="load")
 
-@handle_data_exceptions
-def save_data(file_path: Union[str, Path], data: Any) -> bool:
+@handle_data_exceptions(default_return=False, re_raise=True)
+@log_performance
+def save_data(
+    data: Any,
+    filename: str,
+    storage_dir: Union[str, Path] = None,
+    format: str = "auto",
+    encoding: str = "utf-8",
+    backup: bool = True,
+    atomic: bool = True
+) -> bool:
     """
-    Save data to a JSON file
+    Save data to a file.
     
     Args:
-        file_path: Path to the JSON file
         data: Data to save
-        
+        filename: Name of the file to save to
+        storage_dir: Directory to save the file in (defaults to DEFAULT_STORAGE_DIR)
+        format: File format ('json', 'pickle', 'txt', or 'auto' to detect from extension)
+        encoding: Text encoding for text files
+        backup: Whether to create a backup of existing file
+        atomic: Whether to use atomic write (write to temp file then move)
+    
     Returns:
         True if successful, False otherwise
-        
+    
     Raises:
-        DataError: If there's an error saving the data
+        StorageError: If file cannot be saved
     """
-    file_path = Path(file_path)
+    if storage_dir is None:
+        storage_dir = DEFAULT_STORAGE_DIR
     
-    # Create directory if it doesn't exist
-    file_path.parent.mkdir(parents=True, exist_ok=True)
+    storage_path = Path(storage_dir)
+    storage_path.mkdir(parents=True, exist_ok=True)
+    
+    file_path = storage_path / filename
+    
+    # Auto-detect format from extension
+    if format == "auto":
+        ext = file_path.suffix.lower()
+        if ext == ".json":
+            format = "json"
+        elif ext in [".pkl", ".pickle"]:
+            format = "pickle"
+        else:
+            format = "txt"
+    
+    # Create backup if requested and file exists
+    if backup and file_path.exists():
+        backup_path = file_path.with_suffix(f"{file_path.suffix}.backup")
+        try:
+            shutil.copy2(file_path, backup_path)
+            logger.debug(f"Created backup at {backup_path}")
+        except Exception as e:
+            logger.warning(f"Could not create backup: {e}")
     
     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}")
+        logger.debug(f"Saving data to {file_path} with format {format}")
+        
+        if atomic:
+            # Use atomic write: write to temp file, then move
+            with tempfile.NamedTemporaryFile(mode='w' if format != 'pickle' else 'wb', 
+                                           dir=storage_path, 
+                                           delete=False,
+                                           suffix=f'.tmp_{filename}') as temp_file:
+                temp_path = Path(temp_file.name)
+                
+                if format == "json":
+                    json.dump(data, temp_file, indent=2, ensure_ascii=False)
+                elif format == "pickle":
+                    pickle.dump(data, temp_file)
+                elif format == "txt":
+                    if isinstance(data, str):
+                        temp_file.write(data)
+                    else:
+                        temp_file.write(str(data))
+                else:
+                    raise StorageError(f"Unsupported format: {format}", operation="save")
+            
+            # Atomic move
+            shutil.move(str(temp_path), str(file_path))
+        else:
+            # Direct write
+            if format == "json":
+                with open(file_path, 'w', encoding=encoding) as f:
+                    json.dump(data, f, indent=2, ensure_ascii=False)
+            elif format == "pickle":
+                with open(file_path, 'wb') as f:
+                    pickle.dump(data, f)
+            elif format == "txt":
+                with open(file_path, 'w', encoding=encoding) as f:
+                    if isinstance(data, str):
+                        f.write(data)
+                    else:
+                        f.write(str(data))
+            else:
+                raise StorageError(f"Unsupported format: {format}", operation="save")
+        
+        logger.info(f"Successfully saved data to {file_path}")
         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
+        # Clean up temp file if atomic write failed
+        if atomic and 'temp_path' in locals() and temp_path.exists():
+            try:
+                temp_path.unlink()
+            except:
+                pass
+        raise StorageError(f"Error saving file {file_path}: {e}", operation="save")
 
-@handle_data_exceptions
-def export_data(data: Any, format_type: str = 'json', file_path: Optional[Union[str, Path]] = None) -> Union[str, bool]:
+@handle_data_exceptions(default_return=[])
+def list_files(
+    storage_dir: Union[str, Path] = None,
+    pattern: str = "*",
+    extension: str = None
+) -> List[str]:
     """
-    Export data to various formats
+    List files in storage directory.
     
     Args:
-        data: Data to export
-        format_type: Format to export to ('json', 'csv', 'markdown')
-        file_path: Path to save the exported data (optional)
-        
+        storage_dir: Directory to list files from
+        pattern: Glob pattern to match files
+        extension: File extension to filter by (e.g., '.json')
+    
     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
+        List of filenames
     """
-    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}")
+    if storage_dir is None:
+        storage_dir = DEFAULT_STORAGE_DIR
     
-    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
+    storage_path = Path(storage_dir)
     
-    if file_path is None:
-        return result
+    if not storage_path.exists():
+        return []
     
-    file_path = Path(file_path)
-    file_path.parent.mkdir(parents=True, exist_ok=True)
+    if extension:
+        pattern = f"*{extension}"
     
-    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
-    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
+    files = [f.name for f in storage_path.glob(pattern) if f.is_file()]
+    return sorted(files)
 
-@handle_data_exceptions
-def get_file_path(data_type: str) -> Path:
+@handle_data_exceptions(default_return=False)
+def delete_file(filename: str, storage_dir: Union[str, Path] = None) -> bool:
     """
-    Get the file path for a specific data type
+    Delete a file from storage.
     
     Args:
-        data_type: Type of data (e.g., 'tasks', 'notes', 'goals')
-        
+        filename: Name of the file to delete
+        storage_dir: Directory containing the file
+    
     Returns:
-        Path object for the specified data type
-        
-    Raises:
-        ValidationError: If the data type is invalid
+        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}")
+    if storage_dir is None:
+        storage_dir = DEFAULT_STORAGE_DIR
     
-    return FILE_PATHS[data_type]
+    file_path = Path(storage_dir) / filename
+    
+    try:
+        if file_path.exists():
+            file_path.unlink()
+            logger.info(f"Deleted file {file_path}")
+            return True
+        else:
+            logger.warning(f"File {file_path} does not exist")
+            return False
+    except Exception as e:
+        logger.error(f"Error deleting file {file_path}: {e}")
+        return False
 
-@handle_data_exceptions
-def create_backup(backup_name: str = None) -> Optional[Path]:
+@handle_data_exceptions(default_return={})
+def get_file_info(filename: str, storage_dir: Union[str, Path] = None) -> Dict[str, Any]:
     """
-    Create a backup of all data files
+    Get information about a file.
     
     Args:
-        backup_name: Optional name for the backup file
-        
+        filename: Name of the file
+        storage_dir: Directory containing the file
+    
     Returns:
-        Path to the backup file if successful, None otherwise
+        Dictionary with file information
     """
-    import zipfile
-    from datetime import datetime
+    if storage_dir is None:
+        storage_dir = DEFAULT_STORAGE_DIR
     
-    if backup_name is None:
-        backup_name = f"mona_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    file_path = Path(storage_dir) / filename
     
-    backup_file = EXPORT_DIR / f"{backup_name}.zip"
+    if not file_path.exists():
+        return {}
     
     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
+        stat = file_path.stat()
+        return {
+            "name": filename,
+            "path": str(file_path),
+            "size": stat.st_size,
+            "created": datetime.fromtimestamp(stat.st_ctime).isoformat(),
+            "modified": datetime.fromtimestamp(stat.st_mtime).isoformat(),
+            "extension": file_path.suffix,
+            "exists": True
+        }
     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"Error getting file info for {file_path}: {e}")
+        return {"name": filename, "exists": False, "error": str(e)}
 
-@handle_data_exceptions
-def restore_backup(backup_file: Union[str, Path]) -> bool:
+# Configuration management
+@handle_data_exceptions(default_return={})
+def load_config(config_file: str = "config.json", storage_dir: Union[str, Path] = None) -> Dict[str, Any]:
     """
-    Restore data from a backup file
+    Load configuration from a JSON file.
     
     Args:
-        backup_file: Path to the backup file
-        
+        config_file: Name of the configuration file
+        storage_dir: Directory containing the config file
+    
+    Returns:
+        Configuration dictionary
+    """
+    config = load_data(config_file, storage_dir, format="json", default={})
+    if not isinstance(config, dict):
+        logger.warning(f"Config file {config_file} did not contain a dictionary, using empty config")
+        return {}
+    return config
+
+@handle_data_exceptions(default_return=False)
+def save_config(config: Dict[str, Any], config_file: str = "config.json", storage_dir: Union[str, Path] = None) -> bool:
+    """
+    Save configuration to a JSON file.
+    
+    Args:
+        config: Configuration dictionary to save
+        config_file: Name of the configuration file
+        storage_dir: Directory to save the config file in
+    
     Returns:
         True if successful, False otherwise
     """
-    import zipfile
+    if not isinstance(config, dict):
+        raise ValidationError("Config must be a dictionary")
     
-    backup_file = Path(backup_file)
+    return save_data(config, config_file, storage_dir, format="json")
+
+# Cache management
+class SimpleCache:
+    """Simple in-memory cache with file backing."""
     
-    if not backup_file.exists():
-        logger.error(f"Backup file not found: {backup_file}")
-        raise DataError(f"Backup file not found: {backup_file}")
+    def __init__(self, cache_file: str = "cache.json", storage_dir: Union[str, Path] = None):
+        self.cache_file = cache_file
+        self.storage_dir = storage_dir
+        self._cache = {}
+        self.load_cache()
     
-    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
-    except Exception as e:
-        logger.error(f"Error restoring backup: {str(e)}")
-        raise DataError(f"Error restoring backup", {"original_error": str(e)}) from e
+    def load_cache(self):
+        """Load cache from file."""
+        try:
+            self._cache = load_data(self.cache_file, self.storage_dir, default={})
+            if not isinstance(self._cache, dict):
+                self._cache = {}
+        except Exception as e:
+            logger.warning(f"Could not load cache: {e}")
+            self._cache = {}
+    
+    def save_cache(self):
+        """Save cache to file."""
+        try:
+            save_data(self._cache, self.cache_file, self.storage_dir)
+        except Exception as e:
+            logger.warning(f"Could not save cache: {e}")
+    
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get value from cache."""
+        return self._cache.get(key, default)
+    
+    def set(self, key: str, value: Any, save: bool = True) -> None:
+        """Set value in cache."""
+        self._cache[key] = value
+        if save:
+            self.save_cache()
+    
+    def delete(self, key: str, save: bool = True) -> bool:
+        """Delete key from cache."""
+        if key in self._cache:
+            del self._cache[key]
+            if save:
+                self.save_cache()
+            return True
+        return False
+    
+    def clear(self, save: bool = True) -> None:
+        """Clear all cache."""
+        self._cache.clear()
+        if save:
+            self.save_cache()
+    
+    def keys(self) -> List[str]:
+        """Get all cache keys."""
+        return list(self._cache.keys())
+
+# Global cache instance
+default_cache = SimpleCache()
+
+def get_cache() -> SimpleCache:
+    """Get the default cache instance."""
+    return default_cache