feat：完成大部分数据文件上传解析功能
This commit is contained in:
@@ -0,0 +1,240 @@
|
||||
"""把各种格式读成原始网格（CSV、Excel）或按列的记录（JSON、JSONL），Parquet 直接读成 Arrow 表
|
||||
|
||||
CSV 和 Excel 可能有标题行、多行表头、合计行，所以先读成不带表头的网格，再由 tables.py 识别表区域
|
||||
"""
|
||||
|
||||
import csv
|
||||
import io
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
from python_calamine import (
|
||||
CalamineWorkbook,
|
||||
PasswordError,
|
||||
SheetTypeEnum,
|
||||
SheetVisibleEnum,
|
||||
)
|
||||
|
||||
from app.datafiles.base import MAX_TABLE_ROWS, DataFileError, rows_exceeded
|
||||
from app.schemas.data_file import ParseWarning
|
||||
|
||||
# 原始网格比数据多出的标题、表头、空行，读取时先粗略拦住明显超限的文件，精确的行数在拆表后检查
|
||||
ROW_SLACK = 1000
|
||||
|
||||
# JSON 嵌套对象展开的层数，更深的部分存成 JSON 字符串
|
||||
JSON_FLATTEN_LEVELS = 2
|
||||
|
||||
# JSONL 坏行超过这个比例时整个文件判为失败
|
||||
MAX_BAD_LINE_RATIO = 0.01
|
||||
|
||||
# 单元格最长的字符数；csv 模块默认 128KB，长文本单元格会报错
|
||||
csv.field_size_limit(16 * 1024 * 1024)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Grid:
|
||||
"""不带表头的原始网格；rows[i] 是文件或 sheet 中的第 i + 1 行"""
|
||||
|
||||
sheet: str | None
|
||||
rows: list[list[Any]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Records:
|
||||
"""已有列结构的数据，按列存放"""
|
||||
|
||||
names: list[str]
|
||||
columns: list[list[Any]]
|
||||
warnings: list[ParseWarning] = field(default_factory=list)
|
||||
|
||||
|
||||
def read_delimited(text: str, delimiter: str) -> Grid:
|
||||
rows = []
|
||||
for row in csv.reader(io.StringIO(text, newline=""), delimiter=delimiter):
|
||||
rows.append(row)
|
||||
if len(rows) > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(None)
|
||||
return Grid(None, rows)
|
||||
|
||||
|
||||
def read_excel(path: Path) -> tuple[list[Grid], list[ParseWarning]]:
|
||||
try:
|
||||
book = CalamineWorkbook.from_path(str(path))
|
||||
except PasswordError:
|
||||
raise DataFileError(422, "Excel 文件设置了密码，请去掉密码后重新上传") from None
|
||||
except Exception:
|
||||
raise DataFileError(422, "无法读取 Excel 文件，文件可能已损坏") from None
|
||||
grids, warnings = [], []
|
||||
try:
|
||||
for meta in book.sheets_metadata:
|
||||
# 图表页、宏表等没有单元格数据
|
||||
if meta.typ != SheetTypeEnum.WorkSheet:
|
||||
continue
|
||||
if meta.visible != SheetVisibleEnum.Visible:
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="hidden_sheet_skipped",
|
||||
message=f"跳过了隐藏的 sheet「{meta.name}」",
|
||||
)
|
||||
)
|
||||
continue
|
||||
sheet = book.get_sheet_by_name(meta.name)
|
||||
if sheet.total_height > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(meta.name)
|
||||
# 保留左上角的空白区域，行号和 Excel 里看到的一致
|
||||
rows = sheet.to_python(skip_empty_area=False)
|
||||
if not any(cell != "" for row in rows for cell in row):
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="empty_sheet_skipped",
|
||||
message=f"sheet「{meta.name}」是空的，已跳过",
|
||||
)
|
||||
)
|
||||
continue
|
||||
grids.append(Grid(meta.name, rows))
|
||||
finally:
|
||||
book.close()
|
||||
return grids, warnings
|
||||
|
||||
|
||||
def read_json(value: Any) -> Grid | Records:
|
||||
"""对象数组按记录读取；数组的数组当作网格，第一行通常是表头"""
|
||||
warnings = []
|
||||
if isinstance(value, dict):
|
||||
records, path = _find_records(value)
|
||||
if path:
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="json_path_selected",
|
||||
message=f"使用了 JSON 中「{path}」下的 {len(records)} 条记录",
|
||||
)
|
||||
)
|
||||
else:
|
||||
records = value
|
||||
if not records:
|
||||
raise DataFileError(422, "JSON 里没有找到数据")
|
||||
if all(isinstance(item, list) for item in records):
|
||||
return Grid(None, [[_scalar(cell) for cell in row] for row in records])
|
||||
if all(isinstance(item, dict) for item in records):
|
||||
result = _flatten_records(records)
|
||||
result.warnings = warnings
|
||||
return result
|
||||
if not any(isinstance(item, (dict, list)) for item in records):
|
||||
return Records(["值"], [list(records)], warnings)
|
||||
raise DataFileError(422, "JSON 数组里的元素类型不一致，需要都是对象或都是数组")
|
||||
|
||||
|
||||
def read_jsonl(text: str) -> Records:
|
||||
records, bad, total = [], [], 0
|
||||
for number, line in enumerate(text.splitlines(), start=1):
|
||||
if not line.strip():
|
||||
continue
|
||||
total += 1
|
||||
try:
|
||||
record = json.loads(line)
|
||||
except ValueError:
|
||||
record = None
|
||||
if isinstance(record, dict):
|
||||
records.append(record)
|
||||
else:
|
||||
bad.append(number)
|
||||
if not records:
|
||||
raise DataFileError(422, "JSONL 里没有找到数据")
|
||||
if len(bad) > total * MAX_BAD_LINE_RATIO:
|
||||
raise DataFileError(
|
||||
422, f"JSONL 有 {len(bad)} 行无法解析（例如第 {bad[0]} 行），超过 1%"
|
||||
)
|
||||
result = _flatten_records(records)
|
||||
if bad:
|
||||
result.warnings.append(
|
||||
ParseWarning(
|
||||
code="bad_lines_skipped",
|
||||
message=f"跳过了 {len(bad)} 行无法解析的内容：{_line_list(bad)}",
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def read_parquet(path: Path) -> pa.Table:
|
||||
try:
|
||||
table = pq.read_table(path)
|
||||
except Exception:
|
||||
raise DataFileError(422, "无法读取 Parquet 文件，文件可能已损坏") from None
|
||||
if table.num_rows > MAX_TABLE_ROWS:
|
||||
raise rows_exceeded(None)
|
||||
# struct 展开一层，列名写成 a.b；更深的嵌套由 columns.py 存成 JSON 字符串
|
||||
return table.flatten()
|
||||
|
||||
|
||||
def _find_records(value: dict) -> tuple[list, str | None]:
|
||||
"""在对象的前两层里找最长的对象数组，例如 {"data": [...]}；找不到时按列式或单条记录处理"""
|
||||
candidates = []
|
||||
for key, item in value.items():
|
||||
if _is_record_list(item):
|
||||
candidates.append((len(item), str(key), item))
|
||||
elif isinstance(item, dict):
|
||||
for sub_key, sub in item.items():
|
||||
if _is_record_list(sub):
|
||||
candidates.append((len(sub), f"{key}.{sub_key}", sub))
|
||||
if candidates:
|
||||
_, path, records = max(candidates, key=lambda candidate: candidate[0])
|
||||
return records, path
|
||||
lists = list(value.values())
|
||||
# 列式：{"地区": [...], "销售额": [...]}
|
||||
if (
|
||||
lists
|
||||
and all(isinstance(item, list) for item in lists)
|
||||
and len({len(item) for item in lists}) == 1
|
||||
):
|
||||
return [dict(zip(value, row)) for row in zip(*lists)], None
|
||||
return [value], None
|
||||
|
||||
|
||||
def _is_record_list(value: Any) -> bool:
|
||||
return (
|
||||
isinstance(value, list)
|
||||
and bool(value)
|
||||
and all(isinstance(item, dict) for item in value)
|
||||
)
|
||||
|
||||
|
||||
def _flatten_records(records: list[dict]) -> Records:
|
||||
# 列按首次出现的顺序排列，记录里没有的字段为空
|
||||
index: dict[str, int] = {}
|
||||
columns: list[list[Any]] = []
|
||||
for row_number, record in enumerate(records):
|
||||
flat: dict[str, Any] = {}
|
||||
_flatten(record, "", 0, flat)
|
||||
for name, value in flat.items():
|
||||
if name not in index:
|
||||
index[name] = len(columns)
|
||||
columns.append([None] * row_number)
|
||||
columns[index[name]].append(value)
|
||||
for column in columns:
|
||||
if len(column) == row_number:
|
||||
column.append(None)
|
||||
return Records(list(index), columns)
|
||||
|
||||
|
||||
def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None:
|
||||
for key, value in record.items():
|
||||
name = f"{prefix}.{key}" if prefix else str(key)
|
||||
if isinstance(value, dict) and value and level < JSON_FLATTEN_LEVELS:
|
||||
_flatten(value, name, level + 1, out)
|
||||
else:
|
||||
out[name] = _scalar(value)
|
||||
|
||||
|
||||
def _scalar(value: Any) -> Any:
|
||||
if isinstance(value, (dict, list)):
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
return value
|
||||
|
||||
|
||||
def _line_list(numbers: list[int]) -> str:
|
||||
shown = "、".join(f"第 {number} 行" for number in numbers[:5])
|
||||
return shown + ("等" if len(numbers) > 5 else "")
|
||||
Reference in New Issue
Block a user