138 lines
4.1 KiB
Python
138 lines
4.1 KiB
Python
"""
|
|
CSV handling functionality for Animagine Prompt Node
|
|
"""
|
|
|
|
import pandas as pd
|
|
from typing import Optional, Dict
|
|
import os
|
|
import time
|
|
from .logger import logger
|
|
|
|
class CSVHandler:
|
|
"""
|
|
Handles loading and processing of CSV files containing character data
|
|
"""
|
|
|
|
REQUIRED_COLUMNS = ['GENDER', 'CHARACTER', 'COPYRIGHT']
|
|
|
|
def __init__(self):
|
|
self.current_data = None
|
|
self.cache = {}
|
|
|
|
def validate_csv_structure(self, filepath: str) -> bool:
|
|
"""
|
|
Validates if the CSV file has the required structure
|
|
"""
|
|
try:
|
|
if not os.path.exists(filepath):
|
|
logger.log_error(
|
|
FileNotFoundError(f"CSV file not found: {filepath}"),
|
|
"CSV Validation",
|
|
{"filepath": filepath}
|
|
)
|
|
return False
|
|
|
|
df = pd.read_csv(filepath)
|
|
logger.log_csv_operation(
|
|
"Validation",
|
|
filepath,
|
|
{"columns_found": df.columns.tolist()}
|
|
)
|
|
|
|
missing_columns = [col for col in self.REQUIRED_COLUMNS if col not in df.columns]
|
|
if missing_columns:
|
|
logger.log_error(
|
|
ValueError("Missing required columns"),
|
|
"CSV Validation",
|
|
{"missing_columns": missing_columns}
|
|
)
|
|
return False
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.log_error(e, "CSV Validation", {"filepath": filepath})
|
|
return False
|
|
|
|
def load_csv(self, filepath: str) -> Optional[pd.DataFrame]:
|
|
"""
|
|
Loads CSV file with basic caching
|
|
"""
|
|
try:
|
|
start_time = time.time()
|
|
|
|
if filepath in self.cache:
|
|
logger.log_csv_operation(
|
|
"Cache Hit",
|
|
filepath,
|
|
{"cache_size": len(self.cache)}
|
|
)
|
|
return self.cache[filepath]
|
|
|
|
if not self.validate_csv_structure(filepath):
|
|
return None
|
|
|
|
df = pd.read_csv(filepath)
|
|
self.cache[filepath] = df
|
|
self.current_data = df
|
|
|
|
end_time = time.time()
|
|
logger.log_performance(
|
|
"CSV Load",
|
|
start_time,
|
|
end_time,
|
|
{
|
|
"filepath": filepath,
|
|
"rows": len(df),
|
|
"cached": True
|
|
}
|
|
)
|
|
|
|
return df
|
|
|
|
except Exception as e:
|
|
logger.log_error(e, "CSV Load", {"filepath": filepath})
|
|
return None
|
|
|
|
def get_entry(self, index: int) -> Optional[Dict[str, str]]:
|
|
"""
|
|
Gets a specific entry from the loaded CSV
|
|
"""
|
|
try:
|
|
if self.current_data is None:
|
|
logger.log_error(
|
|
ValueError("No CSV data loaded"),
|
|
"Get Entry",
|
|
{"index": index}
|
|
)
|
|
return None
|
|
|
|
if index >= len(self.current_data):
|
|
logger.log_error(
|
|
IndexError(f"Index {index} out of bounds"),
|
|
"Get Entry",
|
|
{
|
|
"index": index,
|
|
"max_index": len(self.current_data) - 1
|
|
}
|
|
)
|
|
return None
|
|
|
|
row = self.current_data.iloc[index]
|
|
entry = {
|
|
'gender': row['GENDER'],
|
|
'character': row['CHARACTER'],
|
|
'copyright': row['COPYRIGHT']
|
|
}
|
|
|
|
logger.log_csv_operation(
|
|
"Get Entry",
|
|
None,
|
|
{"index": index, "entry": entry}
|
|
)
|
|
|
|
return entry
|
|
|
|
except Exception as e:
|
|
logger.log_error(e, "Get Entry", {"index": index})
|
|
return None |