Update utils/storage.py

utils/storage.py

@@ -1,171 +1,355 @@
+"""
+Storage utilities for the MONA application
+Fixed version with proper error handling and data management
+"""
+
 import json
 import os
 import pickle
+import csv
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
-import pandas as pd
+from typing import Any, Dict, List, Optional, Union
+from datetime import datetime
 
 from utils.error_handling import handle_data_exceptions, DataError, ValidationError
-from utils.logging import get_logger
+from utils.logging import get_logger, log_error, log_info, log_warning
+
+
+class StorageManager:
+    """Manages data storage and retrieval operations"""
+    
+    def __init__(self, base_path: str = "data"):
+        self.base_path = Path(base_path)
+        self.base_path.mkdir(exist_ok=True)
+        self.logger = get_logger(__name__)
+    
+    def get_file_path(self, filename: str, subfolder: str = None) -> Path:
+        """Get the full file path for a given filename"""
+        if subfolder:
+            folder = self.base_path / subfolder
+            folder.mkdir(exist_ok=True)
+            return folder / filename
+        return self.base_path / filename
 
-logger = get_logger(__name__)
 
-# Create data directory if it doesn't exist
-DATA_DIR = Path("data")
-DATA_DIR.mkdir(exist_ok=True)
+# Global storage manager instance
+_storage_manager = StorageManager()
+
 
 @handle_data_exceptions
-def save_data(data: Any, filename: str, format_type: str = "json") -> bool:
+def save_data(data: Any, filename: str, format_type: str = "json", subfolder: str = None) -> bool:
     """
-    Save data to file
+    Save data to file in specified format
     
     Args:
         data: Data to save
-        filename: Name of the file
-        format_type: Format to save ('json', 'pickle', 'csv', 'excel')
-    
+        filename: Name of the file to save
+        format_type: Format to save in ('json', 'pickle', 'csv', 'txt')
+        subfolder: Optional subfolder to save in
+        
     Returns:
-        bool: True if successful, False otherwise
+        bool: True if successful
+        
+    Raises:
+        DataError: If saving fails
+        ValidationError: If parameters are invalid
     """
+    if not filename:
+        raise ValidationError("Filename cannot be empty")
+    
+    if format_type not in ['json', 'pickle', 'csv', 'txt']:
+        raise ValidationError(f"Unsupported format type: {format_type}")
+    
     try:
-        filepath = DATA_DIR / filename
+        file_path = _storage_manager.get_file_path(filename, subfolder)
         
         if format_type == "json":
-            with open(filepath, 'w', encoding='utf-8') as f:
+            with open(file_path, '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:
+            with open(file_path, '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")
+            if not isinstance(data, (list, tuple)):
+                raise ValidationError("CSV format requires list or tuple data")
+            
+            with open(file_path, 'w', newline='', encoding='utf-8') as f:
+                if data and isinstance(data[0], dict):
+                    # Dictionary data
+                    writer = csv.DictWriter(f, fieldnames=data[0].keys())
+                    writer.writeheader()
+                    writer.writerows(data)
+                else:
+                    # List data
+                    writer = csv.writer(f)
+                    writer.writerows(data)
         
-        else:
-            raise DataError(f"Unsupported format: {format_type}")
+        elif format_type == "txt":
+            with open(file_path, 'w', encoding='utf-8') as f:
+                if isinstance(data, str):
+                    f.write(data)
+                else:
+                    f.write(str(data))
         
-        logger.info(f"Data saved successfully to {filepath}")
+        log_info(f"Successfully saved data to {file_path}")
         return True
         
     except Exception as e:
-        logger.error(f"Failed to save data to {filename}: {str(e)}")
-        raise DataError(f"Failed to save data: {str(e)}")
+        raise DataError(f"Failed to save data to {filename}", details={"format": format_type, "error": str(e)})
+
 
 @handle_data_exceptions
-def load_data(filename: str, format_type: str = "json") -> Any:
+def load_data(filename: str, format_type: str = "json", subfolder: str = None, default: Any = None) -> Any:
     """
-    Load data from file
+    Load data from file in specified format
     
     Args:
-        filename: Name of the file
-        format_type: Format to load ('json', 'pickle', 'csv', 'excel')
-    
+        filename: Name of the file to load
+        format_type: Format to load from ('json', 'pickle', 'csv', 'txt')
+        subfolder: Optional subfolder to load from
+        default: Default value if file doesn't exist
+        
     Returns:
-        Any: Loaded data or None if failed
+        Any: Loaded data or default value
+        
+    Raises:
+        DataError: If loading fails
+        ValidationError: If parameters are invalid
     """
+    if not filename:
+        raise ValidationError("Filename cannot be empty")
+    
+    if format_type not in ['json', 'pickle', 'csv', 'txt']:
+        raise ValidationError(f"Unsupported format type: {format_type}")
+    
     try:
-        filepath = DATA_DIR / filename
+        file_path = _storage_manager.get_file_path(filename, subfolder)
         
-        if not filepath.exists():
-            logger.warning(f"File not found: {filepath}")
-            return None
+        if not file_path.exists():
+            if default is not None:
+                log_warning(f"File {file_path} not found, returning default value")
+                return default
+            else:
+                raise DataError(f"File not found: {file_path}")
         
         if format_type == "json":
-            with open(filepath, 'r', encoding='utf-8') as f:
+            with open(file_path, 'r', encoding='utf-8') as f:
                 data = json.load(f)
         
         elif format_type == "pickle":
-            with open(filepath, 'rb') as f:
+            with open(file_path, '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)
+            data = []
+            with open(file_path, 'r', encoding='utf-8') as f:
+                # Try to detect if first row is header
+                sample = f.read(1024)
+                f.seek(0)
+                
+                sniffer = csv.Sniffer()
+                has_header = sniffer.has_header(sample)
+                
+                if has_header:
+                    reader = csv.DictReader(f)
+                    data = list(reader)
+                else:
+                    reader = csv.reader(f)
+                    data = list(reader)
         
-        else:
-            raise DataError(f"Unsupported format: {format_type}")
+        elif format_type == "txt":
+            with open(file_path, 'r', encoding='utf-8') as f:
+                data = f.read()
         
-        logger.info(f"Data loaded successfully from {filepath}")
+        log_info(f"Successfully loaded data from {file_path}")
         return data
         
     except Exception as e:
-        logger.error(f"Failed to load data from {filename}: {str(e)}")
-        raise DataError(f"Failed to load data: {str(e)}")
+        if default is not None:
+            log_warning(f"Failed to load {filename}, returning default: {str(e)}")
+            return default
+        raise DataError(f"Failed to load data from {filename}", details={"format": format_type, "error": str(e)})
+
 
 @handle_data_exceptions
-def delete_data(filename: str) -> bool:
+def delete_file(filename: str, subfolder: str = None) -> bool:
     """
-    Delete data file
+    Delete a file
     
     Args:
         filename: Name of the file to delete
-    
+        subfolder: Optional subfolder
+        
     Returns:
-        bool: True if successful, False otherwise
+        bool: True if successful
+        
+    Raises:
+        DataError: If deletion fails
     """
     try:
-        filepath = DATA_DIR / filename
+        file_path = _storage_manager.get_file_path(filename, subfolder)
         
-        if filepath.exists():
-            filepath.unlink()
-            logger.info(f"File deleted successfully: {filepath}")
+        if file_path.exists():
+            file_path.unlink()
+            log_info(f"Successfully deleted {file_path}")
             return True
         else:
-            logger.warning(f"File not found for deletion: {filepath}")
+            log_warning(f"File {file_path} does not exist")
             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)}")
+        raise DataError(f"Failed to delete {filename}", details={"error": str(e)})
+
 
 @handle_data_exceptions
-def list_data_files() -> list:
+def list_files(subfolder: str = None, extension: str = None) -> List[str]:
     """
-    List all data files in the data directory
+    List files in the storage directory
     
+    Args:
+        subfolder: Optional subfolder to list
+        extension: Optional file extension filter
+        
     Returns:
-        list: List of filenames
+        List[str]: List of filenames
+        
+    Raises:
+        DataError: If listing fails
     """
     try:
-        files = [f.name for f in DATA_DIR.iterdir() if f.is_file()]
-        logger.info(f"Found {len(files)} data files")
-        return files
+        if subfolder:
+            folder_path = _storage_manager.base_path / subfolder
+        else:
+            folder_path = _storage_manager.base_path
+        
+        if not folder_path.exists():
+            return []
+        
+        files = []
+        for file_path in folder_path.iterdir():
+            if file_path.is_file():
+                if extension is None or file_path.suffix.lower() == extension.lower():
+                    files.append(file_path.name)
+        
+        return sorted(files)
         
     except Exception as e:
-        logger.error(f"Failed to list data files: {str(e)}")
-        return []
+        raise DataError("Failed to list files", details={"subfolder": subfolder, "error": str(e)})
+
 
-def get_file_size(filename: str) -> Optional[int]:
+@handle_data_exceptions
+def file_exists(filename: str, subfolder: str = None) -> bool:
     """
-    Get file size in bytes
+    Check if a file exists
     
     Args:
-        filename: Name of the file
+        filename: Name of the file to check
+        subfolder: Optional subfolder
+        
+    Returns:
+        bool: True if file exists
+    """
+    try:
+        file_path = _storage_manager.get_file_path(filename, subfolder)
+        return file_path.exists()
+    except Exception as e:
+        log_error(f"Error checking file existence: {str(e)}", error=e)
+        return False
+
+
+@handle_data_exceptions
+def get_file_info(filename: str, subfolder: str = None) -> Dict:
+    """
+    Get information about a file
     
+    Args:
+        filename: Name of the file
+        subfolder: Optional subfolder
+        
     Returns:
-        int: File size in bytes or None if file doesn't exist
+        Dict: File information including size, modified time, etc.
+        
+    Raises:
+        DataError: If file doesn't exist or info retrieval fails
     """
     try:
-        filepath = DATA_DIR / filename
-        if filepath.exists():
-            return filepath.stat().st_size
-        return None
+        file_path = _storage_manager.get_file_path(filename, subfolder)
+        
+        if not file_path.exists():
+            raise DataError(f"File not found: {filename}")
+        
+        stat = file_path.stat()
+        
+        return {
+            "name": file_path.name,
+            "size": stat.st_size,
+            "modified": datetime.fromtimestamp(stat.st_mtime).isoformat(),
+            "created": datetime.fromtimestamp(stat.st_ctime).isoformat(),
+            "extension": file_path.suffix,
+            "path": str(file_path)
+        }
+        
     except Exception as e:
-        logger.error(f"Failed to get file size for {filename}: {str(e)}")
+        raise DataError(f"Failed to get file info for {filename}", details={"error": str(e)})
+
+
+# Cache management
+class DataCache:
+    """Simple in-memory data cache"""
+    
+    def __init__(self, max_size: int = 100):
+        self.cache: Dict[str, Any] = {}
+        self.access_times: Dict[str, datetime] = {}
+        self.max_size = max_size
+    
+    def get(self, key: str) -> Optional[Any]:
+        """Get value from cache"""
+        if key in self.cache:
+            self.access_times[key] = datetime.now()
+            return self.cache[key]
         return None
+    
+    def set(self, key: str, value: Any) -> None:
+        """Set value in cache"""
+        if len(self.cache) >= self.max_size:
+            # Remove oldest item
+            oldest_key = min(self.access_times.keys(), key=lambda k: self.access_times[k])
+            del self.cache[oldest_key]
+            del self.access_times[oldest_key]
+        
+        self.cache[key] = value
+        self.access_times[key] = datetime.now()
+    
+    def clear(self) -> None:
+        """Clear all cache"""
+        self.cache.clear()
+        self.access_times.clear()
+    
+    def remove(self, key: str) -> bool:
+        """Remove specific key from cache"""
+        if key in self.cache:
+            del self.cache[key]
+            del self.access_times[key]
+            return True
+        return False
+
+
+# Global cache instance
+_data_cache = DataCache()
+
+
+def get_cached_data(key: str) -> Optional[Any]:
+    """Get data from cache"""
+    return _data_cache.get(key)
+
+
+def set_cached_data(key: str, value: Any) -> None:
+    """Set data in cache"""
+    _data_cache.set(key, value)
+
 
-def ensure_data_directory():
-    """Ensure data directory exists"""
-    DATA_DIR.mkdir(exist_ok=True)
-    return DATA_DIR
+def clear_cache() -> None:
+    """Clear all cached data"""
+    _data_cache.clear()