237 lines
8.4 KiB
Python
237 lines
8.4 KiB
Python
"""按文件内容识别格式和文本编码，不看扩展名
|
||
|
||
文本文件整个解码一遍确认编码，但只留开头一段用来判断格式和分隔符；读取时再按确认的编码边读边解码
|
||
"""
|
||
|
||
import codecs
|
||
import csv
|
||
import json
|
||
import re
|
||
import zipfile
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from app.datafiles.base import UNSUPPORTED_FORMAT, DataFileError
|
||
from app.schemas.data_file import DataFormat
|
||
|
||
# xlsx/xlsb 解压后的总大小上限，防止解压炸弹
|
||
MAX_UNZIPPED_BYTES = 1024 * 1024 * 1024
|
||
|
||
# 整体是一个 JSON 值的文件只能整个读进内存解析，峰值内存约为文件大小的 6～7 倍；
|
||
# 这个上限让它和最大的 CSV 一样峰值约 1 GB。JSONL 边读边解析，不受这个限制
|
||
MAX_JSON_BYTES = 150 * 1024 * 1024
|
||
|
||
# 带 BOM 的文本直接按 BOM 解码：编码名、跳过 BOM 后的解码方式、读取时的编码（会自己去掉 BOM）；
|
||
# UTF-32 的 BOM 以 UTF-16 的开头，解出来含 NUL，按二进制拒绝
|
||
BOMS = (
|
||
(codecs.BOM_UTF8, "utf-8", "utf-8", "utf-8-sig"),
|
||
(codecs.BOM_UTF16_LE, "utf-16", "utf-16-le", "utf-16"),
|
||
(codecs.BOM_UTF16_BE, "utf-16", "utf-16-be", "utf-16"),
|
||
)
|
||
|
||
# 没有 BOM 时依次尝试；GB18030 兼容 GBK、GB2312，中文 Windows 上 Excel 另存的 CSV 多是这种
|
||
FALLBACK_ENCODINGS = ("utf-8", "gb18030")
|
||
|
||
DELIMITERS = (",", "\t", ";", "|")
|
||
|
||
# 判断分隔符只看开头这么多行
|
||
SNIFF_LINES = 50
|
||
|
||
# 判断格式和分隔符用的开头文本长度；每次解码读入的字节数
|
||
HEAD_CHARS = 1024 * 1024
|
||
READ_BYTES = 1024 * 1024
|
||
|
||
# 开头像 JSON 才按 JSON 处理：对象的 { 后面只能是 " 或 }，数组的 [ 后面是一个值的开头；
|
||
# CSV 的列名也可能以 [ 开头，例如 [日期],[销售额]
|
||
JSON_START = re.compile(r'\{\s*["}]|\[\s*(?:[\[{"\]\d-]|(?:true|false|null)\s*[,\]])')
|
||
|
||
UNKNOWN_ENCODING = "无法识别文件编码，请另存为 UTF-8 后重新上传"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Detected:
|
||
format: DataFormat
|
||
# 文本格式的编码名（返回给调用方）和读取时传给 open 的编码
|
||
encoding: str | None = None
|
||
codec: str | None = None
|
||
# 文本格式开头的一段
|
||
head: str | None = None
|
||
# CSV/TSV 的分隔符
|
||
delimiter: str | None = None
|
||
# JSON 已解析的内容，避免读取时再解析一遍
|
||
value: Any = None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class _Text:
|
||
encoding: str
|
||
codec: str
|
||
head: str
|
||
# 开头一段之后还有内容
|
||
truncated: bool
|
||
# 整个文件只有空白
|
||
blank: bool
|
||
|
||
|
||
def detect(path: Path) -> Detected:
|
||
with path.open("rb") as file:
|
||
head = file.read(8)
|
||
if head.startswith(b"PK\x03\x04"):
|
||
return Detected(_zip_format(path))
|
||
if head.startswith(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"):
|
||
return Detected("xls")
|
||
if head.startswith(b"PAR1"):
|
||
return Detected("parquet")
|
||
|
||
text = scan_text(path)
|
||
if text.blank:
|
||
raise DataFileError(422, "文件是空的")
|
||
if JSON_START.match(text.head.lstrip()):
|
||
return _json_format(path, text)
|
||
delimiter = sniff_delimiter(text.head)
|
||
return Detected(
|
||
"tsv" if delimiter == "\t" else "csv",
|
||
text.encoding,
|
||
text.codec,
|
||
text.head,
|
||
delimiter,
|
||
)
|
||
|
||
|
||
def _zip_format(path: Path) -> DataFormat:
|
||
try:
|
||
with zipfile.ZipFile(path) as archive:
|
||
infos = archive.infolist()
|
||
except zipfile.BadZipFile:
|
||
raise DataFileError(415, UNSUPPORTED_FORMAT) from None
|
||
if sum(info.file_size for info in infos) > MAX_UNZIPPED_BYTES:
|
||
raise DataFileError(413, "Excel 文件解压后过大")
|
||
names = {info.filename for info in infos}
|
||
if "xl/workbook.xml" in names:
|
||
return "xlsx"
|
||
if "xl/workbook.bin" in names:
|
||
return "xlsb"
|
||
raise DataFileError(415, UNSUPPORTED_FORMAT)
|
||
|
||
|
||
def scan_text(path: Path) -> _Text:
|
||
"""确认文本编码：BOM 优先，否则依次严格尝试 FALLBACK_ENCODINGS"""
|
||
with path.open("rb") as file:
|
||
start = file.read(4)
|
||
for bom, name, decoding, codec in BOMS:
|
||
if start.startswith(bom):
|
||
try:
|
||
return _decode(path, name, decoding, codec, len(bom))
|
||
except UnicodeDecodeError:
|
||
raise DataFileError(422, UNKNOWN_ENCODING) from None
|
||
for encoding in FALLBACK_ENCODINGS:
|
||
try:
|
||
return _decode(path, encoding, encoding, encoding, 0)
|
||
except UnicodeDecodeError:
|
||
continue
|
||
raise DataFileError(422, UNKNOWN_ENCODING)
|
||
|
||
|
||
def _decode(path: Path, name: str, decoding: str, codec: str, skip: int) -> _Text:
|
||
"""严格解码整个文件，留下开头一段；含 NUL 时按二进制文件拒绝"""
|
||
decoder = codecs.getincrementaldecoder(decoding)("strict")
|
||
parts: list[str] = []
|
||
kept, blank, truncated = 0, True, False
|
||
with path.open("rb") as file:
|
||
file.seek(skip)
|
||
while True:
|
||
chunk = file.read(READ_BYTES)
|
||
text = decoder.decode(chunk, final=not chunk)
|
||
# 文本文件不会有 NUL；图片、PDF 等二进制文件偶尔能按 GB18030 解码成功，在这里拦下
|
||
if "\x00" in text:
|
||
raise DataFileError(415, UNSUPPORTED_FORMAT)
|
||
blank = blank and not text.strip()
|
||
if kept < HEAD_CHARS:
|
||
parts.append(text)
|
||
kept += len(text)
|
||
elif text:
|
||
truncated = True
|
||
if not chunk:
|
||
break
|
||
head = "".join(parts)
|
||
if len(head) > HEAD_CHARS:
|
||
head, truncated = head[:HEAD_CHARS], True
|
||
cut = head.rfind("\n")
|
||
if truncated and cut >= 0:
|
||
# 截断处可能在一行中间，丢掉最后半行
|
||
head = head[: cut + 1]
|
||
return _Text(name, codec, head, truncated, blank)
|
||
|
||
|
||
def _json_format(path: Path, text: _Text) -> Detected:
|
||
if _first_line_is_record(text):
|
||
return Detected("jsonl", text.encoding, text.codec, text.head)
|
||
# 整体是一个 JSON 值，只能整个读进来解析
|
||
if path.stat().st_size > MAX_JSON_BYTES:
|
||
raise DataFileError(
|
||
413,
|
||
f"JSON 文件不能超过 {MAX_JSON_BYTES // 1024 // 1024} MB，"
|
||
"请转成 JSONL（每行一个 JSON 对象）后上传",
|
||
)
|
||
content = path.read_bytes().decode(text.codec)
|
||
try:
|
||
return Detected(
|
||
"json", text.encoding, text.codec, text.head, value=json.loads(content)
|
||
)
|
||
except ValueError as error:
|
||
json_error = error
|
||
first = next(line for line in content.splitlines() if line.strip())
|
||
try:
|
||
if isinstance(json.loads(first), dict):
|
||
return Detected("jsonl", text.encoding, text.codec, text.head)
|
||
except ValueError:
|
||
pass
|
||
# 开头像 JSON 却不是合法 JSON，按 JSON 报错比当成 CSV 解析出一堆乱码更有用
|
||
line = getattr(json_error, "lineno", None)
|
||
position = f"第 {line} 行" if line else ""
|
||
raise DataFileError(
|
||
422, f"JSON 格式错误：{position}{getattr(json_error, 'msg', json_error)}"
|
||
)
|
||
|
||
|
||
def _first_line_is_record(text: _Text) -> bool:
|
||
"""第一行是完整的 JSON 对象、后面还有内容：整体不可能是一个 JSON 值，是 JSONL
|
||
|
||
不用整个解析就能判断，大 JSONL 文件不必整个读进内存
|
||
"""
|
||
lines = text.head.splitlines(keepends=True)
|
||
index = next((i for i, line in enumerate(lines) if line.strip()), None)
|
||
if index is None:
|
||
return False
|
||
first, rest = lines[index], lines[index + 1 :]
|
||
if first == first.rstrip("\r\n") and (rest or text.truncated):
|
||
# 第一行在开头一段里没有结束
|
||
return False
|
||
if not (text.truncated or any(line.strip() for line in rest)):
|
||
return False
|
||
try:
|
||
return isinstance(json.loads(first), dict)
|
||
except ValueError:
|
||
return False
|
||
|
||
|
||
def sniff_delimiter(text: str) -> str:
|
||
"""选出让开头各行列数最一致、列数最多的分隔符；都只有一列时按逗号处理"""
|
||
lines = []
|
||
for line in text.splitlines():
|
||
if line.strip():
|
||
lines.append(line)
|
||
if len(lines) == SNIFF_LINES:
|
||
break
|
||
best, best_score = ",", (0.0, 1)
|
||
for delimiter in DELIMITERS:
|
||
widths = [len(row) for row in csv.reader(lines, delimiter=delimiter)]
|
||
width = max(set(widths), key=widths.count)
|
||
if width < 2:
|
||
continue
|
||
score = (widths.count(width) / len(widths), width)
|
||
if score > best_score:
|
||
best, best_score = delimiter, score
|
||
return best
|