122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class UnsupportedWorkbookError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class CorruptWorkbookError(RuntimeError):
|
|
"""The file cannot be opened as a valid Excel workbook.
|
|
|
|
Raised for corrupt/truncated files, wrong signatures and files that
|
|
exceed the resource limits. The message is user-safe: it never contains
|
|
server-side paths or the content-addressed storage filename.
|
|
"""
|
|
|
|
|
|
# Workbook resource guards (defense against decompression bombs and
|
|
# accidentally giant exports). Bank statements are small; these bounds are
|
|
# generous enough for real exports while keeping memory bounded.
|
|
MAX_SHEETS = 50
|
|
MAX_ROWS_PER_SHEET = 200_000
|
|
MAX_COLS_PER_SHEET = 64
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RawSheet:
|
|
name: str
|
|
rows: tuple[tuple[Any, ...], ...]
|
|
|
|
|
|
def read_workbook(path: Path) -> tuple[RawSheet, ...]:
|
|
suffix = path.suffix.lower()
|
|
try:
|
|
if suffix == ".xlsx":
|
|
return _read_xlsx(path)
|
|
if suffix == ".xls":
|
|
return _read_xls(path)
|
|
raise UnsupportedWorkbookError(f"Unsupported workbook type: {suffix}")
|
|
except (UnsupportedWorkbookError, CorruptWorkbookError):
|
|
raise
|
|
except Exception as exc:
|
|
raise CorruptWorkbookError(
|
|
"文件无法读取，可能已损坏或不是有效的 Excel 文件。"
|
|
) from exc
|
|
|
|
|
|
def _read_xlsx(path: Path) -> tuple[RawSheet, ...]:
|
|
try:
|
|
from openpyxl import load_workbook
|
|
except ImportError as exc:
|
|
raise UnsupportedWorkbookError(
|
|
"Reading .xlsx files requires openpyxl. Install requirements.txt."
|
|
) from exc
|
|
|
|
workbook = load_workbook(path, read_only=True, data_only=True)
|
|
try:
|
|
sheets: list[RawSheet] = []
|
|
for worksheet in workbook.worksheets:
|
|
rows: list[tuple[Any, ...]] = []
|
|
for row_index, row in enumerate(worksheet.iter_rows(values_only=True)):
|
|
if row_index >= MAX_ROWS_PER_SHEET:
|
|
raise CorruptWorkbookError(
|
|
f"工作表「{worksheet.title}」超过 {MAX_ROWS_PER_SHEET} 行，已拒绝读取。"
|
|
)
|
|
values = tuple(row)
|
|
if len(values) > MAX_COLS_PER_SHEET:
|
|
raise CorruptWorkbookError(
|
|
f"工作表「{worksheet.title}」列数超过 {MAX_COLS_PER_SHEET}，已拒绝读取。"
|
|
)
|
|
rows.append(values)
|
|
sheets.append(RawSheet(name=worksheet.title, rows=tuple(rows)))
|
|
if len(sheets) > MAX_SHEETS:
|
|
raise CorruptWorkbookError(f"工作簿工作表数量超过 {MAX_SHEETS} 个，已拒绝读取。")
|
|
return tuple(sheets)
|
|
finally:
|
|
workbook.close()
|
|
|
|
|
|
def _read_xls(path: Path) -> tuple[RawSheet, ...]:
|
|
try:
|
|
import xlrd
|
|
except ImportError as exc:
|
|
raise UnsupportedWorkbookError(
|
|
"Reading .xls files requires xlrd. Install requirements.txt."
|
|
) from exc
|
|
|
|
workbook = xlrd.open_workbook(path, on_demand=True)
|
|
sheets: list[RawSheet] = []
|
|
try:
|
|
names = workbook.sheet_names()
|
|
if len(names) > MAX_SHEETS:
|
|
raise CorruptWorkbookError(f"工作簿工作表数量超过 {MAX_SHEETS} 个，已拒绝读取。")
|
|
for sheet_name in names:
|
|
worksheet = workbook.sheet_by_name(sheet_name)
|
|
if worksheet.nrows > MAX_ROWS_PER_SHEET:
|
|
raise CorruptWorkbookError(
|
|
f"工作表「{sheet_name}」超过 {MAX_ROWS_PER_SHEET} 行，已拒绝读取。"
|
|
)
|
|
if worksheet.ncols > MAX_COLS_PER_SHEET:
|
|
raise CorruptWorkbookError(
|
|
f"工作表「{sheet_name}」列数超过 {MAX_COLS_PER_SHEET}，已拒绝读取。"
|
|
)
|
|
rows: list[tuple[Any, ...]] = []
|
|
for row_index in range(worksheet.nrows):
|
|
values: list[Any] = []
|
|
for column_index in range(worksheet.ncols):
|
|
cell = worksheet.cell(row_index, column_index)
|
|
value = cell.value
|
|
if cell.ctype == xlrd.XL_CELL_DATE:
|
|
value = xlrd.xldate_as_datetime(value, workbook.datemode)
|
|
values.append(value)
|
|
rows.append(tuple(values))
|
|
sheets.append(RawSheet(name=sheet_name, rows=tuple(rows)))
|
|
finally:
|
|
workbook.release_resources()
|
|
return tuple(sheets)
|