Fix LAYOUT_TABLE prefix and suffix not being properly added

Author: Belval

textractor/entities/layout.py

@@ -203,9 +203,9 @@ def get_text_and_words(
                 )
             if config.add_prefixes_and_suffixes_as_words:
                 final_words = (
-                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.header_prefix, is_structure=True)] if config.title_prefix else []) + 
+                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.header_prefix, is_structure=True)] if config.header_prefix else []) + 
                     final_words + 
-                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.header_suffix, is_structure=True)] if config.title_suffix else []) 
+                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.header_suffix, is_structure=True)] if config.header_suffix else []) 
                 )
         elif self.layout_type == LAYOUT_SECTION_HEADER:
             final_text, final_words = linearize_children(
@@ -217,9 +217,9 @@ def get_text_and_words(
                 )
             if config.add_prefixes_and_suffixes_as_words:
                 final_words = (
-                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.section_header_prefix, is_structure=True)] if config.title_prefix else []) + 
+                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.section_header_prefix, is_structure=True)] if config.section_header_prefix else []) + 
                     final_words + 
-                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.section_header_suffix, is_structure=True)] if config.title_suffix else []) 
+                    ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.section_header_suffix, is_structure=True)] if config.section_header_suffix else []) 
                 )
         elif self.layout_type == LAYOUT_TEXT:
             final_text, final_words = linearize_children(
@@ -235,6 +235,29 @@ def get_text_and_words(
                 is_layout_table=self.layout_type == LAYOUT_TABLE,
             )
 
+            if config.add_prefixes_and_suffixes_in_text:
+                if self.layout_type == LAYOUT_TABLE:
+                    final_text = (
+                        config.table_layout_prefix + final_text + config.table_layout_suffix
+                    )
+                elif self.layout_type == LAYOUT_KEY_VALUE:
+                    final_text = (
+                        config.table_layout_prefix + final_text + config.table_layout_suffix
+                    )
+            if config.add_prefixes_and_suffixes_as_words:
+                if self.layout_type == LAYOUT_TABLE:
+                    final_words = (
+                        ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.table_layout_prefix, is_structure=True)] if config.table_layout_prefix else []) + 
+                        final_words + 
+                        ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.table_layout_suffix, is_structure=True)] if config.table_layout_suffix else []) 
+                    )
+                elif self.layout_type == LAYOUT_KEY_VALUE:
+                    final_words = (
+                        ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.key_value_layout_prefix, is_structure=True)] if config.key_value_layout_prefix else []) + 
+                        final_words + 
+                        ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.key_value_layout_suffix, is_structure=True)] if config.key_value_layout_suffix else []) 
+                    )
+
         while (
             config.layout_element_separator * (config.max_number_of_consecutive_new_lines + 1) in final_text
         ):

textractor/entities/table.py

@@ -716,6 +716,8 @@ def get_text_and_words(
             w.table_id = str(self.id)
             w.table_bbox = self.bbox
 
+
+        text = (config.table_prefix if config.add_prefixes_and_suffixes_in_text else "")
         # Markdown
         if config.table_linearization_format == "markdown":
             df = self.to_pandas(
@@ -727,12 +729,11 @@ def get_text_and_words(
                 headers = df.columns if has_column else []
             else:
                 headers = df.columns
-            text = df.to_markdown(
+            text += df.to_markdown(
                 index=False, tablefmt=config.table_tabulate_format, headers=headers
             )
         # Plaintext
         else:
-            text = ""
             rows = itertools.groupby(self.table_cells, key=lambda cell: cell.row_index)
             processed_cells = set()
             for _, row in rows:
@@ -764,7 +765,7 @@ def get_text_and_words(
                     text = text[:-1]
                 text += (config.table_row_suffix if config.add_prefixes_and_suffixes_in_text else "")
                 text += "\n"
-
+        text += (config.table_suffix if config.add_prefixes_and_suffixes_in_text else "")
         return text, words
 
     def to_txt(self):

textractor/utils/text_utils.py

@@ -149,9 +149,7 @@ def part_of_same_paragraph(element1, element2, config=config):
         for idx, element in enumerate(sorted_group):
             text_element, words_element = element.get_text_and_words(config)
             if "Table" in element.__class__.__name__ and len(words_element):
-                result += (
-                    config.table_layout_prefix if config.add_prefixes_and_suffixes_in_text else ""
-                ) + text_element
+                result += text_element
                 for w in words_element:
                     added_words.add(w.id)
                 words_output += words_element
@@ -161,11 +159,7 @@ def part_of_same_paragraph(element1, element2, config=config):
                     if prev_element and part_of_same_paragraph(prev_element, element, config) else
                     config.same_layout_element_separator
                 )
-                result += separator + (
-                    (config.key_value_layout_prefix if config.add_prefixes_and_suffixes_in_text else "") +
-                    text_element + 
-                    (config.key_value_layout_suffix if config.add_prefixes_and_suffixes_in_text else "")
-                )
+                result += separator + text_element
                 for w in words_element:
                     added_words.add(w.id)
                 words_output += words_element