fix: fix header levels for DOCX & HTML (#184)
Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
This commit is contained in:
parent
94d0729c50
commit
b9f5c74a7d
@ -180,11 +180,6 @@ class HTMLDocumentBackend(DeclarativeDocumentBackend):
|
|||||||
parent=self.parents[0], label=DocItemLabel.TITLE, text=text
|
parent=self.parents[0], label=DocItemLabel.TITLE, text=text
|
||||||
)
|
)
|
||||||
|
|
||||||
elif hlevel == self.level:
|
|
||||||
self.parents[hlevel] = doc.add_text(
|
|
||||||
parent=self.parents[hlevel - 1], label=label, text=text
|
|
||||||
)
|
|
||||||
|
|
||||||
elif hlevel > self.level:
|
elif hlevel > self.level:
|
||||||
|
|
||||||
# add invisible group
|
# add invisible group
|
||||||
@ -194,10 +189,6 @@ class HTMLDocumentBackend(DeclarativeDocumentBackend):
|
|||||||
label=GroupLabel.SECTION,
|
label=GroupLabel.SECTION,
|
||||||
parent=self.parents[i - 1],
|
parent=self.parents[i - 1],
|
||||||
)
|
)
|
||||||
|
|
||||||
self.parents[hlevel] = doc.add_text(
|
|
||||||
parent=self.parents[hlevel - 1], label=label, text=text
|
|
||||||
)
|
|
||||||
self.level = hlevel
|
self.level = hlevel
|
||||||
|
|
||||||
elif hlevel < self.level:
|
elif hlevel < self.level:
|
||||||
@ -206,12 +197,14 @@ class HTMLDocumentBackend(DeclarativeDocumentBackend):
|
|||||||
for key, val in self.parents.items():
|
for key, val in self.parents.items():
|
||||||
if key > hlevel:
|
if key > hlevel:
|
||||||
self.parents[key] = None
|
self.parents[key] = None
|
||||||
|
|
||||||
self.parents[hlevel] = doc.add_text(
|
|
||||||
parent=self.parents[hlevel - 1], label=label, text=text
|
|
||||||
)
|
|
||||||
self.level = hlevel
|
self.level = hlevel
|
||||||
|
|
||||||
|
self.parents[hlevel] = doc.add_heading(
|
||||||
|
parent=self.parents[hlevel - 1],
|
||||||
|
text=text,
|
||||||
|
level=hlevel,
|
||||||
|
)
|
||||||
|
|
||||||
def handle_paragraph(self, element, idx, doc):
|
def handle_paragraph(self, element, idx, doc):
|
||||||
"""Handles paragraph tags (p)."""
|
"""Handles paragraph tags (p)."""
|
||||||
if element.text is None:
|
if element.text is None:
|
||||||
|
@ -294,13 +294,7 @@ class MsWordDocumentBackend(DeclarativeDocumentBackend):
|
|||||||
level = self.get_level()
|
level = self.get_level()
|
||||||
if isinstance(curr_level, int):
|
if isinstance(curr_level, int):
|
||||||
|
|
||||||
if curr_level == level:
|
if curr_level > level:
|
||||||
|
|
||||||
self.parents[level] = doc.add_heading(
|
|
||||||
parent=self.parents[level - 1], text=text
|
|
||||||
)
|
|
||||||
|
|
||||||
elif curr_level > level:
|
|
||||||
|
|
||||||
# add invisible group
|
# add invisible group
|
||||||
for i in range(level, curr_level):
|
for i in range(level, curr_level):
|
||||||
@ -310,10 +304,6 @@ class MsWordDocumentBackend(DeclarativeDocumentBackend):
|
|||||||
name=f"header-{i}",
|
name=f"header-{i}",
|
||||||
)
|
)
|
||||||
|
|
||||||
self.parents[curr_level] = doc.add_heading(
|
|
||||||
parent=self.parents[curr_level - 1], text=text
|
|
||||||
)
|
|
||||||
|
|
||||||
elif curr_level < level:
|
elif curr_level < level:
|
||||||
|
|
||||||
# remove the tail
|
# remove the tail
|
||||||
@ -322,12 +312,16 @@ class MsWordDocumentBackend(DeclarativeDocumentBackend):
|
|||||||
self.parents[key] = None
|
self.parents[key] = None
|
||||||
|
|
||||||
self.parents[curr_level] = doc.add_heading(
|
self.parents[curr_level] = doc.add_heading(
|
||||||
parent=self.parents[curr_level - 1], text=text
|
parent=self.parents[curr_level - 1],
|
||||||
|
text=text,
|
||||||
|
level=curr_level,
|
||||||
)
|
)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
self.parents[self.level] = doc.add_heading(
|
self.parents[self.level] = doc.add_heading(
|
||||||
parent=self.parents[self.level - 1], text=text
|
parent=self.parents[self.level - 1],
|
||||||
|
text=text,
|
||||||
|
level=1,
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
|
30
tests/test_backend_html.py
Normal file
30
tests/test_backend_html.py
Normal file
@ -0,0 +1,30 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from docling.backend.html_backend import HTMLDocumentBackend
|
||||||
|
from docling.datamodel.base_models import InputFormat
|
||||||
|
from docling.datamodel.document import InputDocument, SectionHeaderItem
|
||||||
|
|
||||||
|
|
||||||
|
def test_heading_levels():
|
||||||
|
in_path = Path("tests/data/wiki_duck.html")
|
||||||
|
in_doc = InputDocument(
|
||||||
|
path_or_stream=in_path,
|
||||||
|
format=InputFormat.HTML,
|
||||||
|
backend=HTMLDocumentBackend,
|
||||||
|
)
|
||||||
|
backend = HTMLDocumentBackend(
|
||||||
|
in_doc=in_doc,
|
||||||
|
path_or_stream=in_path,
|
||||||
|
)
|
||||||
|
doc = backend.convert()
|
||||||
|
|
||||||
|
found_lvl_2 = found_lvl_3 = False
|
||||||
|
for item, _ in doc.iterate_items():
|
||||||
|
if isinstance(item, SectionHeaderItem):
|
||||||
|
if item.text == "Etymology":
|
||||||
|
found_lvl_2 = True
|
||||||
|
assert item.level == 2
|
||||||
|
elif item.text == "Feeding":
|
||||||
|
found_lvl_3 = True
|
||||||
|
assert item.level == 3
|
||||||
|
assert found_lvl_2 and found_lvl_3
|
30
tests/test_backend_msword.py
Normal file
30
tests/test_backend_msword.py
Normal file
@ -0,0 +1,30 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from docling.backend.msword_backend import MsWordDocumentBackend
|
||||||
|
from docling.datamodel.base_models import InputFormat
|
||||||
|
from docling.datamodel.document import InputDocument, SectionHeaderItem
|
||||||
|
|
||||||
|
|
||||||
|
def test_heading_levels():
|
||||||
|
in_path = Path("tests/data/word_sample.docx")
|
||||||
|
in_doc = InputDocument(
|
||||||
|
path_or_stream=in_path,
|
||||||
|
format=InputFormat.DOCX,
|
||||||
|
backend=MsWordDocumentBackend,
|
||||||
|
)
|
||||||
|
backend = MsWordDocumentBackend(
|
||||||
|
in_doc=in_doc,
|
||||||
|
path_or_stream=in_path,
|
||||||
|
)
|
||||||
|
doc = backend.convert()
|
||||||
|
|
||||||
|
found_lvl_1 = found_lvl_2 = False
|
||||||
|
for item, _ in doc.iterate_items():
|
||||||
|
if isinstance(item, SectionHeaderItem):
|
||||||
|
if item.text == "Let\u2019s swim!":
|
||||||
|
found_lvl_1 = True
|
||||||
|
assert item.level == 1
|
||||||
|
elif item.text == "Let\u2019s eat":
|
||||||
|
found_lvl_2 = True
|
||||||
|
assert item.level == 2
|
||||||
|
assert found_lvl_1 and found_lvl_2
|
Loading…
Reference in New Issue
Block a user