235 lines
7.2 KiB
Python
235 lines
7.2 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Excel 解析工具：将 Excel 文件转为 AI 可读的 Markdown 文本
|
||
支持 .xlsx / .xls / .csv
|
||
"""
|
||
import io
|
||
import logging
|
||
from typing import Tuple, List, Dict, Any
|
||
|
||
_logger = logging.getLogger(__name__)
|
||
|
||
# 默认限制
|
||
MAX_ROWS = 500 # 单 Sheet 最大行数（超出的给统计摘要）
|
||
MAX_TOTAL_ROWS = 2000 # 所有 Sheet 总行数上限
|
||
MAX_FILE_SIZE_MB = 20 # 文件大小上限（MB）
|
||
|
||
|
||
def parse_excel_to_markdown(
|
||
file_bytes: bytes,
|
||
filename: str,
|
||
max_rows: int = MAX_ROWS,
|
||
max_total_rows: int = MAX_TOTAL_ROWS,
|
||
) -> Tuple[str, Dict[str, Any]]:
|
||
"""
|
||
将 Excel/CSV 解析为 Markdown 文本，供 AI 阅读。
|
||
|
||
:param file_bytes: 文件二进制内容
|
||
:param filename: 文件名（用于判断类型）
|
||
:param max_rows: 单 Sheet 最大行数
|
||
:param max_total_rows: 所有 Sheet 总行数上限
|
||
:return: (markdown_text, meta_dict)
|
||
meta_dict: {sheets, total_rows, truncated, file_type, error}
|
||
"""
|
||
meta = {
|
||
'sheets': [],
|
||
'total_rows': 0,
|
||
'truncated': False,
|
||
'file_type': '',
|
||
'error': '',
|
||
}
|
||
|
||
try:
|
||
import pandas as pd
|
||
except ImportError:
|
||
meta['error'] = 'pandas 未安装，无法解析 Excel 文件。请运行：pip install pandas openpyxl'
|
||
_logger.error('[ExcelParser] pandas 未安装')
|
||
return '', meta
|
||
|
||
file_type = _detect_file_type(file_bytes, filename)
|
||
meta['file_type'] = file_type
|
||
|
||
try:
|
||
if file_type == 'csv':
|
||
return _parse_csv(file_bytes, max_rows, max_total_rows, meta)
|
||
elif file_type in ('xlsx', 'xls'):
|
||
return _parse_excel(file_bytes, max_rows, max_total_rows, meta)
|
||
else:
|
||
meta['error'] = f'不支持的文件类型：{file_type}'
|
||
return '', meta
|
||
except Exception as e:
|
||
_logger.exception('[ExcelParser] 解析失败: %s', e)
|
||
meta['error'] = f'解析失败：{e}'
|
||
return '', meta
|
||
|
||
|
||
def _detect_file_type(file_bytes: bytes, filename: str) -> str:
|
||
"""根据文件内容和扩展名判断类型"""
|
||
# 1. 优先看扩展名
|
||
lower_name = filename.lower()
|
||
if lower_name.endswith('.csv'):
|
||
return 'csv'
|
||
if lower_name.endswith('.xlsx'):
|
||
return 'xlsx'
|
||
if lower_name.endswith('.xls'):
|
||
return 'xls'
|
||
|
||
# 2. 看魔数（文件头）
|
||
if file_bytes[:2] == b'PK': # ZIP (xlsx)
|
||
return 'xlsx'
|
||
if file_bytes[:2] == b'\xd0\xcf': # Old Excel (xls)
|
||
return 'xls'
|
||
# 尝试当作 CSV
|
||
try:
|
||
file_bytes.decode('utf-8-sig')
|
||
return 'csv'
|
||
except:
|
||
pass
|
||
return 'unknown'
|
||
|
||
|
||
def _parse_csv(bytes_data: bytes, max_rows: int, max_total: int, meta: dict) -> Tuple[str, dict]:
|
||
"""解析 CSV 文件"""
|
||
import pandas as pd
|
||
|
||
# 尝试多种编码
|
||
encoding = _detect_encoding(bytes_data)
|
||
df = pd.read_csv(io.BytesIO(bytes_data), encoding=encoding, nrows=max_total + 1)
|
||
|
||
return _df_to_markdown(df, 'CSV', max_rows, max_total, meta)
|
||
|
||
|
||
def _parse_excel(bytes_data: bytes, max_rows: int, max_total: int, meta: dict) -> Tuple[str, dict]:
|
||
"""解析 Excel 文件（支持多 Sheet）"""
|
||
import pandas as pd
|
||
|
||
xl = pd.ExcelFile(io.BytesIO(bytes_data))
|
||
sheet_names = xl.sheet_names
|
||
|
||
meta['sheets'] = sheet_names
|
||
|
||
parts = []
|
||
total_rows = 0
|
||
|
||
for sheet_name in sheet_names:
|
||
df = pd.read_excel(xl, sheet_name=sheet_name, nrows=max_total + 1 - total_rows)
|
||
|
||
# 检查是否超出总行数
|
||
if total_rows + len(df) > max_total:
|
||
df = df.head(max_total - total_rows)
|
||
meta['truncated'] = True
|
||
_logger.warning('[ExcelParser] 总行数超过 %s，已截断', max_total)
|
||
|
||
sheet_md, sheet_meta = _df_to_markdown(df, sheet_name, max_rows, max_total, meta, is_single_sheet=(len(sheet_names) == 1))
|
||
parts.append(sheet_md)
|
||
total_rows += len(df)
|
||
|
||
if meta.get('truncated'):
|
||
break
|
||
|
||
meta['total_rows'] = total_rows
|
||
return '\n\n'.join(parts), meta
|
||
|
||
|
||
def _df_to_markdown(
|
||
df,
|
||
sheet_name: str,
|
||
max_rows: int,
|
||
max_total: int,
|
||
meta: dict,
|
||
is_single_sheet: bool = False,
|
||
) -> Tuple[str, dict]:
|
||
"""
|
||
将 DataFrame 转为 Markdown。
|
||
策略：
|
||
- 行数 <= max_rows：完整输出
|
||
- 行数 > max_rows：输出统计摘要 + 前 max_rows 行
|
||
"""
|
||
row_count = len(df)
|
||
col_count = len(df.columns)
|
||
|
||
if not is_single_sheet:
|
||
parts = [f'## Sheet: {sheet_name}']
|
||
else:
|
||
parts = []
|
||
|
||
parts.append(f'（{row_count} 行 × {col_count} 列）\n')
|
||
|
||
if row_count == 0:
|
||
parts.append('（空表）')
|
||
return '\n'.join(parts), meta
|
||
|
||
# 统计摘要（数值列）
|
||
numeric_cols = df.select_dtypes(include=['number']).columns.tolist()
|
||
if numeric_cols:
|
||
parts.append('【数值列统计】')
|
||
for col in numeric_cols[:10]: # 最多10个数值列
|
||
series = df[col].dropna()
|
||
if len(series) > 0:
|
||
parts.append(f' - {col}：最小={series.min():.2f}，最大={series.max():.2f}，平均={series.mean():.2f}，合计={series.sum():.2f}')
|
||
parts.append('')
|
||
|
||
# 分类列（唯一值数量）
|
||
cat_cols = [c for c in df.columns if c not in numeric_cols]
|
||
if cat_cols:
|
||
parts.append('【分类列概况】')
|
||
for col in cat_cols[:5]: # 最多5个分类列
|
||
unique_count = df[col].nunique()
|
||
top_val = df[col].mode().iloc[0] if len(df[col].mode()) > 0 else '（无）'
|
||
parts.append(f' - {col}：共 {unique_count} 种值，最常见「{top_val}」')
|
||
parts.append('')
|
||
|
||
# 数据内容
|
||
if row_count <= max_rows:
|
||
parts.append('【完整数据】')
|
||
parts.append(df.to_markdown(index=False))
|
||
else:
|
||
parts.append(f'【前 {max_rows} 行数据】（共 {row_count} 行，已截断）')
|
||
parts.append(df.head(max_rows).to_markdown(index=False))
|
||
meta['truncated'] = True
|
||
|
||
return '\n'.join(parts), meta
|
||
|
||
|
||
def _detect_encoding(bytes_data: bytes) -> str:
|
||
"""检测 CSV 文件编码"""
|
||
for enc in ['utf-8-sig', 'utf-8', 'gbk', 'gb2312', 'iso-8859-1']:
|
||
try:
|
||
bytes_data.decode(enc)
|
||
return enc
|
||
except:
|
||
continue
|
||
return 'utf-8'
|
||
|
||
|
||
def validate_excel_file(file_bytes: bytes, filename: str) -> Dict[str, Any]:
|
||
"""
|
||
验证 Excel 文件是否合法、大小是否超限。
|
||
返回：{valid, error, file_type, size_mb}
|
||
"""
|
||
size_mb = len(file_bytes) / 1024 / 1024
|
||
|
||
if size_mb > MAX_FILE_SIZE_MB:
|
||
return {
|
||
'valid': False,
|
||
'error': f'文件过大（{size_mb:.1f}MB），上限 {MAX_FILE_SIZE_MB}MB',
|
||
'file_type': '',
|
||
'size_mb': round(size_mb, 1),
|
||
}
|
||
|
||
file_type = _detect_file_type(file_bytes, filename)
|
||
if file_type == 'unknown':
|
||
return {
|
||
'valid': False,
|
||
'error': f'不支持的文件格式，仅支持 .xlsx / .xls / .csv',
|
||
'file_type': file_type,
|
||
'size_mb': round(size_mb, 1),
|
||
}
|
||
|
||
return {
|
||
'valid': True,
|
||
'error': '',
|
||
'file_type': file_type,
|
||
'size_mb': round(size_mb, 1),
|
||
}
|