File size: 2,647 Bytes
ea2348b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | # 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.
"""
@staticmethod
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
@staticmethod
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
|