149 lines
5.1 KiB
Python
149 lines
5.1 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
|
|
|
|
# 带 BOM 的文本直接按 BOM 解码；UTF-32 的 BOM 以 UTF-16 的开头，解出来含 NUL，按二进制拒绝
|
|
BOMS = (
|
|
(codecs.BOM_UTF8, "utf-8", "utf-8"),
|
|
(codecs.BOM_UTF16_LE, "utf-16-le", "utf-16"),
|
|
(codecs.BOM_UTF16_BE, "utf-16-be", "utf-16"),
|
|
)
|
|
|
|
# 没有 BOM 时依次尝试；GB18030 兼容 GBK、GB2312，中文 Windows 上 Excel 另存的 CSV 多是这种
|
|
FALLBACK_ENCODINGS = ("utf-8", "gb18030")
|
|
|
|
DELIMITERS = (",", "\t", ";", "|")
|
|
|
|
# 判断分隔符只看开头这么多行
|
|
SNIFF_LINES = 50
|
|
|
|
# 开头像 JSON 才按 JSON 处理：对象的 { 后面只能是 " 或 }，数组的 [ 后面是一个值的开头；
|
|
# CSV 的列名也可能以 [ 开头，例如 [日期],[销售额]
|
|
JSON_START = re.compile(r'\{\s*["}]|\[\s*(?:[\[{"\]\d-]|(?:true|false|null)\s*[,\]])')
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Detected:
|
|
format: DataFormat
|
|
encoding: str | None = None
|
|
# 文本格式解码后的内容
|
|
text: str | None = None
|
|
# CSV/TSV 的分隔符
|
|
delimiter: str | None = None
|
|
# JSON 已解析的内容，避免读取时再解析一遍
|
|
value: Any = None
|
|
|
|
|
|
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, encoding = decode_text(path.read_bytes())
|
|
stripped = text.lstrip()
|
|
if not stripped:
|
|
raise DataFileError(422, "文件是空的")
|
|
if JSON_START.match(stripped):
|
|
return _json_format(text, encoding)
|
|
delimiter = sniff_delimiter(text)
|
|
return Detected("tsv" if delimiter == "\t" else "csv", encoding, text, 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 decode_text(data: bytes) -> tuple[str, str]:
|
|
"""返回解码后的文本和编码名"""
|
|
for bom, codec, name in BOMS:
|
|
if data.startswith(bom):
|
|
try:
|
|
text = data[len(bom) :].decode(codec)
|
|
except UnicodeDecodeError:
|
|
raise DataFileError(
|
|
422, "无法识别文件编码，请另存为 UTF-8 后重新上传"
|
|
) from None
|
|
return _reject_binary(text), name
|
|
for encoding in FALLBACK_ENCODINGS:
|
|
try:
|
|
return _reject_binary(data.decode(encoding)), encoding
|
|
except UnicodeDecodeError:
|
|
continue
|
|
raise DataFileError(422, "无法识别文件编码，请另存为 UTF-8 后重新上传")
|
|
|
|
|
|
def _reject_binary(text: str) -> str:
|
|
# 文本文件不会有 NUL；图片、PDF 等二进制文件偶尔能按 GB18030 解码成功，在这里拦下
|
|
if "\x00" in text:
|
|
raise DataFileError(415, UNSUPPORTED_FORMAT)
|
|
return text
|
|
|
|
|
|
def _json_format(text: str, encoding: str) -> Detected:
|
|
try:
|
|
return Detected("json", encoding, text, value=json.loads(text))
|
|
except ValueError as error:
|
|
json_error = error
|
|
first = next(line for line in text.splitlines() if line.strip())
|
|
try:
|
|
if isinstance(json.loads(first), dict):
|
|
return Detected("jsonl", encoding, text)
|
|
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 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
|