236 lines
8.2 KiB
Python
236 lines
8.2 KiB
Python
"""把各种格式读成原始网格（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, brief, rows_exceeded
|
|
from app.datafiles.columns import as_text
|
|
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)} 行无法解析的内容：{brief([f'第 {n} 行' for n in 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:
|
|
# 嵌套的对象、数组存成 JSON 字符串，其余原样交给类型推断
|
|
return as_text(value) if isinstance(value, (dict, list)) else value
|