a more concise and effective extractor for excel and csv files.

pull/20625/head
haiyangpengai 12 months ago
parent b33f8b47ca
commit ce0910d2fe

@ -396,20 +396,42 @@ def _extract_text_from_csv(file_content: bytes) -> str:
if not rows: if not rows:
return "" return ""
# Combine multi-line text in the header row
header_row = [cell.replace("\n", " ").replace("\r", "") for cell in rows[0]]
# Create Markdown table # Create Markdown table
markdown_table = "| " + " | ".join(rows[0]) + " |\n" markdown_table = "| " + " | ".join(header_row) + " |\n"
markdown_table += "| " + " | ".join(["---"] * len(rows[0])) + " |\n" markdown_table += "| " + " | ".join(["-" * len(col) for col in rows[0]]) + " |\n"
# Process each data row and combine multi-line text in each cell
for row in rows[1:]: for row in rows[1:]:
markdown_table += "| " + " | ".join(row) + " |\n" processed_row = [cell.replace("\n", " ").replace("\r", "") for cell in row]
markdown_table += "| " + " | ".join(processed_row) + " |\n"
return markdown_table.strip() return markdown_table
except Exception as e: except Exception as e:
raise TextExtractionError(f"Failed to extract text from CSV: {str(e)}") from e raise TextExtractionError(f"Failed to extract text from CSV: {str(e)}") from e
def _extract_text_from_excel(file_content: bytes) -> str: def _extract_text_from_excel(file_content: bytes) -> str:
"""Extract text from an Excel file using pandas.""" """Extract text from an Excel file using pandas."""
def _construct_markdown_table(df: pd.DataFrame) -> str:
"""Manually construct a Markdown table from a DataFrame."""
# Construct the header row
header_row = "| " + " | ".join(df.columns) + " |"
# Construct the separator row
separator_row = "| " + " | ".join(["-" * len(col) for col in df.columns]) + " |"
# Construct the data rows
data_rows = []
for _, row in df.iterrows():
data_row = "| " + " | ".join(map(str, row)) + " |"
data_rows.append(data_row)
# Combine all rows into a single string
markdown_table = "\n".join([header_row, separator_row] + data_rows)
return markdown_table
try: try:
excel_file = pd.ExcelFile(io.BytesIO(file_content)) excel_file = pd.ExcelFile(io.BytesIO(file_content))
markdown_table = "" markdown_table = ""
@ -417,8 +439,15 @@ def _extract_text_from_excel(file_content: bytes) -> str:
try: try:
df = excel_file.parse(sheet_name=sheet_name) df = excel_file.parse(sheet_name=sheet_name)
df.dropna(how="all", inplace=True) df.dropna(how="all", inplace=True)
# Create Markdown table two times to separate tables with a newline
markdown_table += df.to_markdown(index=False, floatfmt="") + "\n\n" # Combine multi-line text in each cell into a single line
df = df.applymap(lambda x: ' '.join(str(x).splitlines()) if isinstance(x, str) else x)
# Combine multi-line text in column names into a single line
df.columns = [' '.join(col.splitlines()) for col in df.columns]
# Manually construct the Markdown table
markdown_table += _construct_markdown_table(df) + "\n\n"
except Exception as e: except Exception as e:
continue continue
return markdown_table return markdown_table
@ -587,4 +616,4 @@ def _extract_text_from_properties(file_content: bytes) -> str:
return "\n".join(result) return "\n".join(result)
except Exception as e: except Exception as e:
raise TextExtractionError(f"Failed to extract text from properties file: {str(e)}") from e raise TextExtractionError(f"Failed to extract text from properties file: {str(e)}") from e
Loading…
Cancel
Save