muhammadjalaludin
Add dataset and file handler classes for transaction data processing; include Pydantic models for API validation
ea2348b | # ta_ml/data/file_handler.py | |
| """ | |
| A utility class for handling file operations, such as parsing uploaded files. | |
| """ | |
| from io import BytesIO | |
| from typing import IO, Literal, Optional | |
| import pandas as pd | |
| from fastapi import UploadFile | |
| class FileHandlerError(Exception): | |
| """Custom exception for file handling errors.""" | |
| pass | |
| class FileHandler: | |
| """ | |
| A class providing static methods to handle file reading and writing. | |
| """ | |
| def read_to_dataframe(file: UploadFile) -> pd.DataFrame: | |
| """ | |
| Reads an uploaded file (CSV or Excel) into a pandas DataFrame. | |
| Args: | |
| file (UploadFile): The file uploaded via a FastAPI endpoint. | |
| Returns: | |
| pd.DataFrame: The content of the file as a pandas DataFrame. | |
| Raises: | |
| FileHandlerError: If the file format is unsupported or if there's | |
| a reading error. | |
| """ | |
| if file.filename.endswith(".csv"): | |
| parser = pd.read_csv | |
| elif file.filename.endswith((".xls", ".xlsx")): | |
| parser = pd.read_excel | |
| else: | |
| raise FileHandlerError( | |
| "Unsupported file format. Please use CSV or Excel (.xls, .xlsx)." | |
| ) | |
| try: | |
| return parser(file.file) | |
| except Exception as e: | |
| raise FileHandlerError(f"Error reading file: {e}") from e | |
| def export_dataframe( | |
| df: pd.DataFrame, | |
| format: Literal["csv", "xlsx"] | |
| ) -> tuple[BytesIO, str, str]: | |
| """ | |
| Exports a pandas DataFrame to an in-memory file (CSV or XLSX). | |
| Args: | |
| df (pd.DataFrame): The DataFrame to export. | |
| format (Literal["csv", "xlsx"]): The desired output format. | |
| Returns: | |
| A tuple containing: | |
| - BytesIO: The in-memory buffer containing the file data. | |
| - str: The appropriate MIME type for the format. | |
| - str: The suggested filename for the download. | |
| """ | |
| buffer = BytesIO() | |
| if format == "csv": | |
| df.to_csv(buffer, index=False) | |
| media_type = "text/csv" | |
| filename = "export.csv" | |
| elif format == "xlsx": | |
| with pd.ExcelWriter(buffer, engine="openpyxl") as writer: | |
| df.to_excel(writer, index=False, sheet_name="Export") | |
| media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" | |
| filename = "export.xlsx" | |
| else: | |
| raise ValueError("Unsupported export format specified.") | |
| buffer.seek(0) | |
| return buffer, media_type, filename | |