mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-11-13 06:40:51 +08:00
Lukas Buck
7aed80c043
Fix Excel parser failing on one-sheet workbooks (required len(df) > 1, but pandas always returns a dict) — sheet name was lost and file read twice. Now accepts any dict size.
363 lines
11 KiB
Python
363 lines
11 KiB
Python
#
|
||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
#
|
||
|
||
import importlib.util
|
||
import os
|
||
import sys
|
||
from io import BytesIO
|
||
from unittest import mock
|
||
|
||
import pytest
|
||
|
||
# Import RAGFlowExcelParser directly by file path to avoid triggering
|
||
# deepdoc/parser/__init__.py and rag.nlp, which pull in heavy dependencies.
|
||
for _m in ["rag.nlp", "rag.utils", "rag.utils.lazy_image"]:
|
||
if _m not in sys.modules:
|
||
sys.modules[_m] = mock.MagicMock()
|
||
|
||
|
||
def _find_project_root(marker="pyproject.toml"):
|
||
d = os.path.dirname(os.path.abspath(__file__))
|
||
while d != os.path.dirname(d):
|
||
if os.path.exists(os.path.join(d, marker)):
|
||
return d
|
||
d = os.path.dirname(d)
|
||
return None
|
||
|
||
|
||
_PROJECT_ROOT = _find_project_root()
|
||
_spec = importlib.util.spec_from_file_location(
|
||
"deepdoc.parser.excel_parser",
|
||
os.path.join(_PROJECT_ROOT, "deepdoc", "parser", "excel_parser.py"),
|
||
)
|
||
_mod = importlib.util.module_from_spec(_spec)
|
||
sys.modules["deepdoc.parser.excel_parser"] = _mod
|
||
_spec.loader.exec_module(_mod)
|
||
|
||
RAGFlowExcelParser = _mod.RAGFlowExcelParser
|
||
|
||
|
||
def _test_find_codec(binary):
|
||
for encoding in ("utf-8-sig", "gb18030"):
|
||
try:
|
||
binary.decode(encoding)
|
||
return encoding
|
||
except UnicodeDecodeError:
|
||
continue
|
||
raise UnicodeDecodeError("test", binary, 0, 1, "undecodable")
|
||
|
||
|
||
_mod.find_codec = _test_find_codec
|
||
_mod.decode_text = lambda binary, document_type="text": (binary.decode(_test_find_codec(binary)), _test_find_codec(binary))
|
||
|
||
|
||
def _make_xlsx(n_data_rows):
|
||
from openpyxl import Workbook
|
||
|
||
wb = Workbook()
|
||
ws = wb.active
|
||
ws.append(["H1", "H2"])
|
||
for i in range(n_data_rows):
|
||
ws.append([f"a{i}", f"b{i}"])
|
||
buf = BytesIO()
|
||
wb.save(buf)
|
||
buf.seek(0)
|
||
return buf.read()
|
||
|
||
|
||
def _chunk_has_no_data_cells(chunk):
|
||
return "<td>" not in chunk and "<td></td>" not in chunk
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_exact_multiple_does_not_emit_header_only_chunk():
|
||
# 12 data rows with chunk_rows=12 (the value rag/app/naive.py uses).
|
||
chunks = RAGFlowExcelParser().html(_make_xlsx(12), chunk_rows=12)
|
||
assert len(chunks) == 1
|
||
assert all(not _chunk_has_no_data_cells(c[0]) for c in chunks)
|
||
assert chunks[0][1][0] == 0 and chunks[0][1][1] == 2 and chunks[0][1][2] == 13
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_multiple_of_chunk_rows_splits_without_spurious_chunk():
|
||
# 24 data rows with chunk_rows=12 -> exactly 2 data chunks, no trailing header-only chunk.
|
||
chunks = RAGFlowExcelParser().html(_make_xlsx(24), chunk_rows=12)
|
||
assert len(chunks) == 2
|
||
assert all(not _chunk_has_no_data_cells(c[0]) for c in chunks)
|
||
assert chunks[0][1][1] == 2 and chunks[0][1][2] == 13
|
||
assert chunks[1][1][1] == 14 and chunks[1][1][2] == 25
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_html_col_max_uses_widest_row():
|
||
from openpyxl import Workbook
|
||
|
||
wb = Workbook()
|
||
ws = wb.active
|
||
ws.append(["H1"])
|
||
ws.append(["a", "b", "c"])
|
||
buf = BytesIO()
|
||
wb.save(buf)
|
||
buf.seek(0)
|
||
chunks = RAGFlowExcelParser().html(buf.read(), chunk_rows=12)
|
||
assert chunks[0][1][3] == 1 and chunks[0][1][4] == 3
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_non_multiple_unchanged():
|
||
# 13 data rows with chunk_rows=12 -> 2 chunks (12 + 1).
|
||
chunks = RAGFlowExcelParser().html(_make_xlsx(13), chunk_rows=12)
|
||
assert len(chunks) == 2
|
||
assert all(not _chunk_has_no_data_cells(c[0]) for c in chunks)
|
||
|
||
|
||
def _make_xlsx_with_values(header, row):
|
||
from openpyxl import Workbook
|
||
|
||
wb = Workbook()
|
||
ws = wb.active
|
||
ws.append(header)
|
||
ws.append(row)
|
||
buf = BytesIO()
|
||
wb.save(buf)
|
||
buf.seek(0)
|
||
return buf.read()
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_call_keeps_zero_valued_cells():
|
||
# __call__ produces the text used for indexing. A numeric 0 (and 0.0 / False)
|
||
# is real data, not an empty cell, so it must survive. The header is only
|
||
# emitted alongside a kept value, so a dropped 0 also loses its "stock" label.
|
||
lines = RAGFlowExcelParser()(_make_xlsx_with_values(["name", "stock"], ["widget", 0]))
|
||
joined = " ".join(text for text, _ in lines)
|
||
assert "stock" in joined and "0" in joined, lines
|
||
assert lines[0][1][1] == 2
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_call_omits_a_blank_header_instead_of_labelling_it_none():
|
||
# A blank header cell has no label to give. str(None) is "None", which is
|
||
# truthy, so it defeats the separator guard and the literal token "None" is
|
||
# indexed and shown alongside the value it pretends to describe.
|
||
lines = RAGFlowExcelParser()(_make_xlsx_with_values(["name", None, "city"], ["widget", "note-1", "paris"]))
|
||
joined = " ".join(text for text, _ in lines)
|
||
assert joined == "name：widget; note-1; city：paris", lines
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_call_keeps_a_zero_header():
|
||
# Guards the tempting shorter fix, str(ti[i].value or ""), which would drop
|
||
# a numeric 0 header the same way it drops a None one. 0 is a real label.
|
||
lines = RAGFlowExcelParser()(_make_xlsx_with_values(["name", 0], ["widget", "note-1"]))
|
||
joined = " ".join(text for text, _ in lines)
|
||
assert "0：note-1" in joined, lines
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_call_skips_truly_empty_cells():
|
||
# None / empty-string cells carry no value and should still be skipped.
|
||
lines = RAGFlowExcelParser()(_make_xlsx_with_values(["name", "note"], ["widget", None]))
|
||
joined = " ".join(text for text, _ in lines)
|
||
assert "note" not in joined, lines
|
||
|
||
|
||
@pytest.mark.p2
|
||
@pytest.mark.parametrize("encoding", ["utf-8", "utf-8-sig", "gbk", "gb18030"])
|
||
def test_csv_encoding_is_detected_without_losing_unicode(encoding):
|
||
binary = "项目名称,备注\n核心系统重构,包含GBK语料导入测试\n".encode(encoding)
|
||
lines = RAGFlowExcelParser()(binary)
|
||
joined = " ".join(text for text, _ in lines)
|
||
assert "项目名称" in joined
|
||
assert "包含GBK语料导入测试" in joined
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_csv_accepts_gb18030_characters_not_supported_by_gbk():
|
||
binary = "项目名称,备注\n扩展字符,𠀀\n".encode("gb18030")
|
||
lines = RAGFlowExcelParser()(binary)
|
||
assert "备注：𠀀" in lines[0][0]
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_corrupt_csv_encoding_still_fails():
|
||
with pytest.raises(Exception, match="Failed to parse CSV"):
|
||
RAGFlowExcelParser()(b"\xff\xfe\xfa\xfb")
|
||
|
||
|
||
def _make_two_sheet_xlsx():
|
||
from openpyxl import Workbook
|
||
|
||
wb = Workbook()
|
||
ws1 = wb.active
|
||
ws1.title = "s1"
|
||
ws1.append(["H"])
|
||
ws1.append(["a"])
|
||
ws2 = wb.create_sheet("s2")
|
||
ws2.append(["H"])
|
||
ws2.append(["b"])
|
||
buf = BytesIO()
|
||
wb.save(buf)
|
||
buf.seek(0)
|
||
return buf.read()
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_call_emits_zero_based_sheet_index():
|
||
# Flow JSON stores this index unchanged; add_positions later does pn+1 so
|
||
# sheet 2 is stored as 2 and the preview selects xs.datas[1].
|
||
lines = RAGFlowExcelParser()(_make_two_sheet_xlsx())
|
||
sheets = {pos[0] for _, pos in lines}
|
||
assert sheets == {0, 1}
|
||
second = next(pos for _, pos in lines if pos[0] == 1)
|
||
assert second == (1, 2, 2, 1, 1)
|
||
|
||
|
||
def _make_large_used_range_xlsx(first_data_row, last_data_row, inflate_row=12000):
|
||
from openpyxl import Workbook
|
||
from openpyxl.styles import Font
|
||
|
||
wb = Workbook()
|
||
ws = wb.active
|
||
ws.cell(row=first_data_row, column=1, value="question")
|
||
ws.cell(row=first_data_row, column=2, value="answer")
|
||
for r in range(first_data_row + 1, last_data_row + 1):
|
||
ws.cell(row=r, column=1, value=f"Q{r}")
|
||
ws.cell(row=r, column=2, value=f"A{r}")
|
||
ws.cell(row=inflate_row, column=1).font = Font(bold=True)
|
||
buf = BytesIO()
|
||
wb.save(buf)
|
||
buf.seek(0)
|
||
return buf.read()
|
||
|
||
|
||
def _make_gap_xlsx(gap_rows, inflate_row=12000):
|
||
from openpyxl import Workbook
|
||
|
||
wb = Workbook()
|
||
ws = wb.active
|
||
ws["A1"] = "Name"
|
||
ws["B1"] = "City"
|
||
for r in range(2, 202):
|
||
ws.cell(row=r, column=1, value=f"n{r}")
|
||
ws.cell(row=r, column=2, value=f"c{r}")
|
||
second_block_start = 202 + gap_rows
|
||
for r in range(second_block_start, second_block_start + 200):
|
||
ws.cell(row=r, column=1, value=f"n{r}")
|
||
ws.cell(row=r, column=2, value=f"c{r}")
|
||
ws.cell(row=inflate_row, column=3, value=" ")
|
||
buf = BytesIO()
|
||
wb.save(buf)
|
||
buf.seek(0)
|
||
return buf.read()
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_large_used_range_with_blank_preamble_parses_all_rows():
|
||
# Regression for #19236: data below row 100 must not be treated as empty.
|
||
xlsx = _make_large_used_range_xlsx(150, 350)
|
||
rows = RAGFlowExcelParser()(xlsx)
|
||
assert len(rows) == 201
|
||
assert "Q350" in rows[-1][0]
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_large_used_range_with_blank_gap_parses_both_blocks():
|
||
# Regression for #19185: a blank gap must not truncate rows after it.
|
||
xlsx = _make_gap_xlsx(gap_rows=600)
|
||
rows = RAGFlowExcelParser()(xlsx)
|
||
assert len(rows) == 400
|
||
assert "n201" in rows[199][0]
|
||
assert "n802" in rows[200][0]
|
||
|
||
|
||
@pytest.mark.p2
|
||
def test_row_number_includes_rows_after_blank_gap():
|
||
xlsx = _make_gap_xlsx(gap_rows=600)
|
||
total = RAGFlowExcelParser.row_number("test.xlsx", xlsx)
|
||
assert total >= 1001
|
||
|
||
|
||
def _one_sheet_xlsx(sheet_title="Revenue"):
|
||
from openpyxl import Workbook
|
||
|
||
wb = Workbook()
|
||
ws = wb.active
|
||
ws.title = sheet_title
|
||
ws.append(["Quarter", "Amount"])
|
||
ws.append(["Q1", 10])
|
||
buffer = BytesIO()
|
||
wb.save(buffer)
|
||
buffer.seek(0)
|
||
return buffer
|
||
|
||
|
||
def test_dataframe_to_workbook_accepts_a_single_sheet_dict():
|
||
"""`pd.read_excel(sheet_name=None)` returns a dict whatever the sheet count."""
|
||
import pandas as pd
|
||
|
||
frame = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]})
|
||
|
||
workbook = RAGFlowExcelParser._dataframe_to_workbook({"Revenue": frame})
|
||
|
||
assert workbook.sheetnames == ["Revenue"]
|
||
|
||
|
||
def test_pandas_fallback_keeps_the_sheet_name_of_a_one_sheet_workbook(monkeypatch):
|
||
"""The fallback taken when openpyxl cannot read the file.
|
||
|
||
The sheet name ends up in the chunk as the table caption, so losing it to the
|
||
placeholder "Data" degrades every chunk the workbook produces.
|
||
"""
|
||
|
||
def _openpyxl_cannot_read(*args, **kwargs):
|
||
raise ValueError("openpyxl cannot read this file")
|
||
|
||
monkeypatch.setattr(_mod, "load_workbook", _openpyxl_cannot_read)
|
||
|
||
workbook = RAGFlowExcelParser._load_excel_to_workbook(_one_sheet_xlsx())
|
||
|
||
assert workbook.sheetnames == ["Revenue"]
|
||
assert list(workbook["Revenue"].iter_rows(values_only=True)) == [
|
||
("Quarter", "Amount"),
|
||
("Q1", 10),
|
||
]
|
||
|
||
|
||
def test_pandas_fallback_still_keeps_every_sheet_of_a_multi_sheet_workbook(monkeypatch):
|
||
from openpyxl import Workbook
|
||
|
||
wb = Workbook()
|
||
first = wb.active
|
||
first.title = "Revenue"
|
||
first.append(["Quarter", "Amount"])
|
||
first.append(["Q1", 10])
|
||
second = wb.create_sheet(title="Notes")
|
||
second.append(["Comment"])
|
||
second.append(["all good"])
|
||
buffer = BytesIO()
|
||
wb.save(buffer)
|
||
buffer.seek(0)
|
||
|
||
def _openpyxl_cannot_read(*args, **kwargs):
|
||
raise ValueError("openpyxl cannot read this file")
|
||
|
||
monkeypatch.setattr(_mod, "load_workbook", _openpyxl_cannot_read)
|
||
|
||
workbook = RAGFlowExcelParser._load_excel_to_workbook(buffer)
|
||
|
||
assert workbook.sheetnames == ["Revenue", "Notes"]
|