Alex Galarce
c84decca8d
Feature/rework table handling * test: add cases for configuration boundary * test: add table post-processing tests (cases to ensure preservation of post-table text and metadata during table operations) * test: enhance table handling tests for small and large tables, ensuring proper suffixing and splitting behavior * Fix and clarify intent in comments * fix: enhance handling of post-table text when combining tables * refactor: remove commented-out code in clean_tables function for clarity * test: update post-table text preservation assertions and add debug tracing * refactor: remove unused remove_table_from_page function and its test * refactor: enhance combine method documentation and clarify split_by_size_with_smart_tables logic * Merge remote-tracking branch 'origin/main' into feature/rework-table-handling * refactor: reduce large table threshold from 1000 to 50 rows for better table splitting * Add simple splitting * refactor: uncomment dynamic answer retrieval in run_one_to_n_prompts for improved functionality * Merge remote-tracking branch 'origin/main' into feature/rework-table-handling * refactor: remove unused table handling functions and related constants for cleaner code * refactor: remove unused test functions in preparation for new tests * refactor: streamline continuation table handling by utilizing existing flags * refactor: update table combination logic to include all tables from the main page * refactor: enhance table handling to automatically resolve column mismatches, respect row limits, and improve unit/integration test coverage * rename clean_tables_simple to clean_tables * Update docstrings * refactor: rename table handling functions * Fix indentation and add tests * isort, black * Merge remote-tracking branch 'origin/main' into feature/rework-table-handling Approved-by: Katon Minhas
860 lines
34 KiB
Python
860 lines
34 KiB
Python
import logging
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import pandas as pd
|
|
from sympy import content
|
|
|
|
from src.investment.table_funcs import (Table, combine_continuous_tables,
|
|
convert_str_to_dataframe,
|
|
get_table_info, recreate_page,
|
|
split_large_table_with_headers,
|
|
split_tables_by_rows)
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
|
|
class TestTableFunctions(unittest.TestCase):
|
|
def setUp(self):
|
|
# New list-of-rows format
|
|
self.sample_text_dict_new = {
|
|
"1": "Some metadata\n-------Table Start--------\nTable 1\n[['A', 'B'], ['1', '3'], ['2', '4']]\n-------Table End--------",
|
|
"2": "More metadata\n-------Table Start--------\n\n[['100', '200'], ['5', '7'], ['6', '8']]\n-------Table End--------",
|
|
}
|
|
# Old dict format for backward compatibility testing
|
|
self.sample_text_dict_old = {
|
|
"1": "Some metadata\n-------Table Start--------\nTable 1\n{'A': ['1', '2'], 'B': ['3', '4']}\n-------Table End--------",
|
|
"2": "More metadata\n-------Table Start--------\n\n{'A': ['5', '6'], 'B': ['7', '8']}\n-------Table End--------",
|
|
}
|
|
# Multiple tables on one page
|
|
self.multi_table_text = {
|
|
"1": "Page metadata\n-------Table Start--------\nFirst Table\n[['Col1', 'Col2'], ['A', 'B'], ['C', 'D']]\n-------Table End--------\n\nSome text\n-------Table Start--------\nSecond Table\n[['X', 'Y'], ['1', '2']]\n-------Table End--------"
|
|
}
|
|
# Long table for splitting tests
|
|
self.long_table_text = {
|
|
"1": "Long table metadata\n-------Table Start--------\nLong Table\n[['Name', 'Value'], ['Row1', '1'], ['Row2', '2'], ['Row3', '3'], ['Row4', '4'], ['Row5', '5']]\n-------Table End--------"
|
|
}
|
|
|
|
def test_convert_str_to_dataframe_new_format(self):
|
|
"""Test conversion from new list-of-rows format."""
|
|
table_str = "[['A', 'B'], ['1', '3'], ['2', '4']]"
|
|
df = convert_str_to_dataframe(table_str)
|
|
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
self.assertEqual(df.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertEqual(df.columns[0], "A")
|
|
self.assertEqual(df.iloc[0, 0], "1")
|
|
self.assertEqual(df.iloc[1, 0], "2")
|
|
self.assertEqual(df.iloc[0, 1], "3")
|
|
self.assertTrue(
|
|
df.dtypes.eq("object").all()
|
|
) # All columns should be of type str (object)
|
|
|
|
def test_convert_str_to_dataframe_old_format(self):
|
|
"""Test conversion from old dict format with deprecation warning."""
|
|
table_str = "{'A': ['1', '2'], 'B': ['3', '4']}"
|
|
with self.assertLogs(level="INFO") as log:
|
|
df = convert_str_to_dataframe(table_str)
|
|
|
|
# Should warn about old format
|
|
self.assertTrue(
|
|
any("old table (dict) format detected" in message for message in log.output)
|
|
)
|
|
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
self.assertEqual(df.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertTrue(
|
|
df.dtypes.eq("object").all()
|
|
) # All columns should be of type str (object)
|
|
self.assertEqual(df.iloc[0, 0], "1")
|
|
self.assertEqual(df.iloc[0, 1], "3")
|
|
self.assertEqual(df.iloc[1, 0], "2")
|
|
self.assertEqual(df.iloc[1, 1], "4")
|
|
|
|
def test_convert_str_to_dataframe_invalid_format(self):
|
|
"""Test handling of invalid formats."""
|
|
invalid_str = "not a valid table format"
|
|
df = convert_str_to_dataframe(invalid_str)
|
|
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
self.assertTrue(df.empty)
|
|
|
|
def test_table_init_new_format(self):
|
|
"""Test Table initialization with new list-of-rows format."""
|
|
table_data = [["A", "B"], ["1", "3"], ["2", "4"]]
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
self.assertEqual(table.page_num, "1")
|
|
self.assertEqual(table.metadata, "metadata")
|
|
self.assertEqual(table.header, "Table 1")
|
|
self.assertEqual(table.table.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertEqual(list(table.table.columns), ["A", "B"])
|
|
self.assertEqual(table.table.iloc[0, 0], "1")
|
|
|
|
def test_table_init_dataframe(self):
|
|
"""Test Table initialization with a DataFrame."""
|
|
table_data = pd.DataFrame([["A", "B"], ["1", "3"]], dtype=str)
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
self.assertEqual(table.table.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertTrue(
|
|
table.table.dtypes.eq("object").all()
|
|
) # All columns should be of type str (object)
|
|
|
|
def test_table_init_column_mismatch_padding(self):
|
|
"""Test table initialization handles column count mismatches by padding."""
|
|
# Headers have 3 columns, but data rows have different counts
|
|
table_data = [
|
|
["A", "B", "C"], # 3 columns (header)
|
|
["1", "2"], # 2 columns - should be padded
|
|
["3", "4", "5", "6"], # 4 columns - determines max_cols = 4
|
|
["7"], # 1 column - should be padded
|
|
]
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Test Table",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
# Should have 4 columns (max from all rows) and 3 data rows
|
|
self.assertEqual(table.table.shape, (3, 4))
|
|
self.assertEqual(
|
|
list(table.table.columns), ["A", "B", "C", ""]
|
|
) # Header padded too
|
|
|
|
# Check first row (padded to 4 columns)
|
|
self.assertEqual(table.table.iloc[0, 0], "1")
|
|
self.assertEqual(table.table.iloc[0, 1], "2")
|
|
self.assertEqual(table.table.iloc[0, 2], "") # Padded
|
|
self.assertEqual(table.table.iloc[0, 3], "") # Padded
|
|
|
|
# Check second row (kept at 4 columns)
|
|
self.assertEqual(table.table.iloc[1, 0], "3")
|
|
self.assertEqual(table.table.iloc[1, 1], "4")
|
|
self.assertEqual(table.table.iloc[1, 2], "5")
|
|
self.assertEqual(table.table.iloc[1, 3], "6") # 4th column preserved
|
|
|
|
# Check third row (heavily padded)
|
|
self.assertEqual(table.table.iloc[2, 0], "7")
|
|
self.assertEqual(table.table.iloc[2, 1], "") # Padded
|
|
self.assertEqual(table.table.iloc[2, 2], "") # Padded
|
|
self.assertEqual(table.table.iloc[2, 3], "") # Padded
|
|
|
|
def test_table_init_continuation_column_handling(self):
|
|
"""Test continuation table handles column mismatches with generic column names."""
|
|
# Continuation table - no header row, just data with varying lengths
|
|
table_data = [
|
|
["1", "2"], # 2 columns
|
|
["3", "4", "5"], # 3 columns
|
|
["6"], # 1 column
|
|
]
|
|
table = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
)
|
|
|
|
# Should use max column count (3) and generate generic column names
|
|
self.assertEqual(table.table.shape, (3, 3))
|
|
self.assertEqual(list(table.table.columns), ["col_0", "col_1", "col_2"])
|
|
|
|
# Check padding/truncation
|
|
self.assertEqual(table.table.iloc[0, 0], "1")
|
|
self.assertEqual(table.table.iloc[0, 1], "2")
|
|
self.assertEqual(table.table.iloc[0, 2], "") # Padded
|
|
|
|
self.assertEqual(table.table.iloc[1, 0], "3")
|
|
self.assertEqual(table.table.iloc[1, 1], "4")
|
|
self.assertEqual(table.table.iloc[1, 2], "5")
|
|
|
|
self.assertEqual(table.table.iloc[2, 0], "6")
|
|
self.assertEqual(table.table.iloc[2, 1], "") # Padded
|
|
self.assertEqual(table.table.iloc[2, 2], "") # Padded
|
|
|
|
def test_combine_continuous_tables_simple(self):
|
|
"""Test the simplified continuous table combination logic."""
|
|
# Create test data with continuation table
|
|
text_dict = {
|
|
"1": "Metadata 1\n-------Table Start--------\nMain Table\n[['Name', 'Value'], ['A', '1'], ['B', '2']]\n-------Table End--------",
|
|
"2": "-------Table Start--------\n\n[['C', '3'], ['D', '4']]\n-------Table End--------", # Continuation (empty header)
|
|
"3": "Metadata 3\n-------Table Start--------\nAnother Table\n[['X', 'Y'], ['Z', '5']]\n-------Table End--------",
|
|
}
|
|
|
|
# Debug: Let's see what get_table_info thinks about these tables
|
|
table_dict = get_table_info(text_dict)
|
|
print(f"Table dict keys: {list(table_dict.keys())}")
|
|
for page_num, tables in table_dict.items():
|
|
for i, table in enumerate(tables):
|
|
print(
|
|
f"Page {page_num}, Table {i}: is_continuation={table.is_continuation}, header='{table.header}'"
|
|
)
|
|
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
print(f"Result keys: {list(result.keys())}")
|
|
|
|
# Page 2 should be removed (combined into page 1)
|
|
self.assertNotIn("2", result)
|
|
self.assertIn("1", result)
|
|
self.assertIn("3", result)
|
|
|
|
# Page 1 should now contain the combined table
|
|
combined_table_dict = get_table_info({"1": result["1"]})
|
|
combined_table = combined_table_dict["1"][0]
|
|
|
|
# Should have 4 rows total (2 from main + 2 from continuation)
|
|
self.assertEqual(combined_table.table.shape, (4, 2))
|
|
self.assertEqual(combined_table.table.iloc[0, 0], "A") # First table data
|
|
self.assertEqual(combined_table.table.iloc[2, 0], "C") # Continuation data
|
|
self.assertEqual(combined_table.table.iloc[3, 0], "D") # Continuation data
|
|
|
|
def test_combine_continuous_tables_multiple_tables_per_page(self):
|
|
"""Test that continuation tables combine with the LAST table on previous page, not the first."""
|
|
# Page 1 has TWO tables: Table A (complete) and Table B (continues on page 2)
|
|
text_dict = {
|
|
"1": (
|
|
"Metadata 1\n"
|
|
"-------Table Start--------\nTable A (Complete)\n[['Col1', 'Col2'], ['A1', 'A2'], ['A3', 'A4']]\n-------Table End--------\n"
|
|
"Some text between tables\n"
|
|
"-------Table Start--------\nTable B (Continues)\n[['Name', 'Value'], ['B1', '1'], ['B2', '2']]\n-------Table End--------"
|
|
),
|
|
"2": (
|
|
"-------Table Start--------\n\n[['B3', '3'], ['B4', '4']]\n-------Table End--------" # Continuation of Table B
|
|
),
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict)
|
|
|
|
# Verify setup: Page 1 should have 2 tables, Page 2 should have 1 continuation table
|
|
self.assertEqual(len(table_dict["1"]), 2)
|
|
self.assertEqual(len(table_dict["2"]), 1)
|
|
self.assertFalse(
|
|
table_dict["1"][0].is_continuation
|
|
) # Table A - not continuation
|
|
self.assertFalse(
|
|
table_dict["1"][1].is_continuation
|
|
) # Table B - not continuation
|
|
self.assertTrue(
|
|
table_dict["2"][0].is_continuation
|
|
) # Table B continuation - is continuation
|
|
|
|
# Get original row counts before combining
|
|
table_a_original_rows = table_dict["1"][0].table.shape[0] # Should be 2 rows
|
|
table_b_original_rows = table_dict["1"][1].table.shape[0] # Should be 2 rows
|
|
continuation_rows = table_dict["2"][0].table.shape[0] # Should be 2 rows
|
|
|
|
# Run the combination
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
|
|
# Page 2 should be removed (combined into page 1)
|
|
self.assertNotIn("2", result)
|
|
self.assertIn("1", result)
|
|
|
|
# Parse the combined result to verify tables
|
|
combined_table_dict = get_table_info({"1": result["1"]})
|
|
combined_tables = combined_table_dict["1"]
|
|
|
|
# Should still have 2 tables on page 1
|
|
self.assertEqual(len(combined_tables), 2)
|
|
|
|
# Table A (first table) should be unchanged
|
|
table_a_after = combined_tables[0]
|
|
self.assertEqual(
|
|
table_a_after.table.shape[0], table_a_original_rows
|
|
) # Same row count
|
|
self.assertEqual(table_a_after.header, "Table A (Complete)")
|
|
self.assertEqual(
|
|
table_a_after.table.iloc[0, 0], "A1"
|
|
) # Original data preserved
|
|
|
|
# Table B (second table) should now include the continuation data
|
|
table_b_after = combined_tables[1]
|
|
expected_combined_rows = table_b_original_rows + continuation_rows # 2 + 2 = 4
|
|
self.assertEqual(table_b_after.table.shape[0], expected_combined_rows)
|
|
self.assertEqual(table_b_after.header, "Table B (Continues)")
|
|
|
|
# Verify Table B has both original and continuation data
|
|
self.assertEqual(table_b_after.table.iloc[0, 0], "B1") # Original data
|
|
self.assertEqual(table_b_after.table.iloc[1, 0], "B2") # Original data
|
|
self.assertEqual(table_b_after.table.iloc[2, 0], "B3") # Continuation data
|
|
self.assertEqual(table_b_after.table.iloc[3, 0], "B4") # Continuation data
|
|
|
|
def test_combine_continuous_tables_wrong_combination_scenario(self):
|
|
"""Test edge case: what if we accidentally used [0] instead of [-1]?"""
|
|
# This test documents what WOULD happen with the wrong indexing
|
|
# Page 1: Table A + Table B (continues)
|
|
# Page 2: Table B continuation
|
|
# If we used [0], continuation would wrongly combine with Table A
|
|
|
|
text_dict = {
|
|
"1": (
|
|
"-------Table Start--------\nTable A\n[['X', 'Y'], ['X1', 'Y1']]\n-------Table End--------\n"
|
|
"-------Table Start--------\nTable B\n[['Name', 'Value'], ['B1', '1']]\n-------Table End--------"
|
|
),
|
|
"2": "-------Table Start--------\n\n[['B2', '2']]\n-------Table End--------", # Should combine with Table B, not Table A
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict)
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
|
|
# Get the result and verify correct combination
|
|
combined_table_dict = get_table_info({"1": result["1"]})
|
|
|
|
# Table A (index 0) should have 1 row still - NOT combined with continuation
|
|
table_a = combined_table_dict["1"][0]
|
|
self.assertEqual(table_a.table.shape[0], 1) # Still just 1 row
|
|
self.assertEqual(table_a.table.iloc[0, 0], "X1") # Original data only
|
|
|
|
# Table B (index 1, which is -1) should have 2 rows - correctly combined
|
|
table_b = combined_table_dict["1"][1]
|
|
self.assertEqual(table_b.table.shape[0], 2) # 1 original + 1 continuation = 2
|
|
self.assertEqual(table_b.table.iloc[0, 0], "B1") # Original
|
|
self.assertEqual(table_b.table.iloc[1, 0], "B2") # Continuation
|
|
|
|
def test_table_combine_same_header(self):
|
|
"""Test combining two tables with the same header."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["Column 1", "Column 2"], ["A", "B"], ["1", "2"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 2",
|
|
table_data=[["Column 1", "Column 2"], ["C", "D"], ["3", "4"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
original_shape = table1.table.shape
|
|
table1.combine(table2)
|
|
self.assertEqual(
|
|
table1.table.shape[0], original_shape[0] + table2.table.shape[0]
|
|
) # Rows should be combined
|
|
self.assertEqual(table1.table.shape, (4, 2)) # 4 rows (2 + 2 rows), 2 columns
|
|
self.assertEqual(table1.table.iloc[0, 0], "A")
|
|
self.assertEqual(
|
|
table1.table.iloc[2, 0], "C"
|
|
) # Check if second table's data is appended correctly
|
|
self.assertEqual(table1.table.iloc[3, 1], "4") # Last row of second table
|
|
|
|
def test_table_combine_different_headers(self):
|
|
"""Test combining two tables with different headers (the second table is assumed to be a continuation in this case).
|
|
We will pass in table_data as STRING for both we use convert_str_to_dataframe and
|
|
check that `is_continuation` is getting instantiated correctly.
|
|
"""
|
|
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data="[['Column 1', 'Column 2'], ['A', 'B'], ['1', '2']]",
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
# Continuation table with empty header
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="",
|
|
table_data="[['C', 'D'], ['3', '4'], ['E', 'F']]", # Pass as string!
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True, # Mark as continuation
|
|
)
|
|
|
|
original_shape = table1.table.shape
|
|
table1.combine(table2)
|
|
|
|
print(table1.table)
|
|
self.assertEqual(
|
|
table1.table.shape[0], original_shape[0] + table2.table.shape[0]
|
|
) # Rows should be combined
|
|
self.assertEqual(table1.table.shape, (5, 2)) # 5 rows (2 + 3 rows), 2 columns
|
|
self.assertEqual(table1.table.iloc[0, 0], "A")
|
|
self.assertEqual(
|
|
table1.table.iloc[2, 0], "C"
|
|
) # Check if second table's data is appended correctly
|
|
self.assertEqual(table1.table.iloc[3, 1], "4") # Last row of second table
|
|
self.assertEqual(
|
|
list(table1.table.columns), ["Column 1", "Column 2"]
|
|
) # Should keep the first table's header
|
|
|
|
def test_table_combine_different_columns(self):
|
|
"""Test that combining tables with different column counts handles mismatches gracefully."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["A", "B"], ["1", "2"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 2",
|
|
table_data=[["C", "D", "E"], ["3", "4", "5"]], # 3 columns vs 2 columns
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
# Should not raise an error anymore - should handle column mismatch
|
|
table1.combine(table2)
|
|
|
|
# Verify the result
|
|
# Should have 2 rows total (1 from each table)
|
|
self.assertEqual(table1.table.shape[0], 2)
|
|
|
|
# Should have 2 columns (from the main table)
|
|
self.assertEqual(table1.table.shape[1], 2)
|
|
|
|
# Column names should match the main table
|
|
self.assertEqual(list(table1.table.columns), ["A", "B"])
|
|
|
|
# First row should be from table1
|
|
self.assertEqual(table1.table.iloc[0, 0], "1")
|
|
self.assertEqual(table1.table.iloc[0, 1], "2")
|
|
|
|
# Second row should be from table2, but truncated to 2 columns
|
|
self.assertEqual(table1.table.iloc[1, 0], "3")
|
|
self.assertEqual(table1.table.iloc[1, 1], "4")
|
|
# The "5" from column "E" should be truncated/lost
|
|
|
|
def test_table_combine_column_padding(self):
|
|
"""Test that combining tables pads shorter tables with empty columns."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["A", "B", "C"], ["1", "2", "3"]], # 3 columns
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 2",
|
|
table_data=[["X"], ["4"]], # Only 1 column
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Should have 2 rows total
|
|
self.assertEqual(table1.table.shape[0], 2)
|
|
|
|
# Should have 3 columns (from the main table)
|
|
self.assertEqual(table1.table.shape[1], 3)
|
|
|
|
# Column names should match the main table
|
|
self.assertEqual(list(table1.table.columns), ["A", "B", "C"])
|
|
|
|
# First row should be unchanged
|
|
self.assertEqual(table1.table.iloc[0, 0], "1")
|
|
self.assertEqual(table1.table.iloc[0, 1], "2")
|
|
self.assertEqual(table1.table.iloc[0, 2], "3")
|
|
|
|
# Second row should be from table2, padded with empty strings
|
|
self.assertEqual(table1.table.iloc[1, 0], "4")
|
|
self.assertEqual(table1.table.iloc[1, 1], "") # Padded
|
|
self.assertEqual(table1.table.iloc[1, 2], "") # Padded
|
|
|
|
def test_table_combine_empty_tables(self):
|
|
"""Test combining with empty tables."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="meta",
|
|
header="Table 1",
|
|
table_data=[["A", "B"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
empty_table = Table(
|
|
page_num="2",
|
|
metadata="meta",
|
|
header="",
|
|
table_data=[],
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
)
|
|
|
|
original_shape = table1.table.shape
|
|
table1.combine(empty_table)
|
|
self.assertEqual(table1.table.shape, original_shape) # Should remain unchanged
|
|
|
|
def test_convert_str_to_dataframe_malformed_data(self):
|
|
"""Test handling of malformed table data."""
|
|
malformed_inputs = [
|
|
"[['A', 'B'], ['1']]", # Inconsistent row lengths
|
|
"[['A'], [], ['B']]", # Empty rows
|
|
"[[]]", # Empty nested list
|
|
"['not', 'nested']", # Not nested list
|
|
]
|
|
|
|
for malformed_input in malformed_inputs:
|
|
with self.subTest(input=malformed_input):
|
|
df = convert_str_to_dataframe(malformed_input)
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
|
|
def test_table_post_table_text_preservation(self):
|
|
"""Test that post-table text is properly preserved through operations."""
|
|
text_dict = {
|
|
"1": "Metadata\n-------Table Start--------\nTable\n[['A', 'B'], ['1', '2']]\n-------Table End--------\nImportant footnote\nSignature line"
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict)
|
|
table = table_dict["1"][0]
|
|
|
|
# Should capture post-table text
|
|
self.assertIn("Important footnote", table.post_table_text)
|
|
self.assertIn("Signature line", table.post_table_text)
|
|
|
|
def test_split_large_table_preserves_metadata_and_post_text(self):
|
|
"""Test that table splitting preserves all important information."""
|
|
table = Table(
|
|
page_num="31",
|
|
metadata="EXHIBIT F",
|
|
header="Rate Table",
|
|
table_data=[["HCPC", "Rate"]]
|
|
+ [[f"T{i:04d}", f"${i}.00"] for i in range(15)],
|
|
start_index=100,
|
|
end_index=2000,
|
|
post_table_text="*Rates effective 2024",
|
|
)
|
|
|
|
with patch("src.config.TABLE_ROW_LIMIT", 5):
|
|
chunks = split_large_table_with_headers(table)
|
|
|
|
# All chunks should have metadata and header
|
|
for chunk in chunks:
|
|
self.assertIn("EXHIBIT F", chunk)
|
|
self.assertIn("Rate Table", chunk)
|
|
|
|
# Only last chunk should have post-table text
|
|
self.assertIn("*Rates effective 2024", chunks[-1])
|
|
for chunk in chunks[:-1]:
|
|
self.assertNotIn("*Rates effective 2024", chunk)
|
|
|
|
def test_table_split(self):
|
|
"""Test splitting a table."""
|
|
long_data = [["Name", "Value"]] + [
|
|
[f"Row{i}", str(i)] for i in range(1, 6)
|
|
] # This table will have 5 rows
|
|
long_table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Long Table",
|
|
table_data=long_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
split_tables = long_table.split(row_limit=2)
|
|
|
|
self.assertEqual(
|
|
len(split_tables), 3
|
|
) # Should split into 3 tables (5 rows split by 2 = 3)
|
|
self.assertEqual(
|
|
split_tables[0].table.shape, (2, 2)
|
|
) # First split should have 2 rows
|
|
self.assertEqual(
|
|
split_tables[1].table.shape, (2, 2)
|
|
) # Second split should have 2 rows
|
|
self.assertEqual(
|
|
split_tables[2].table.shape, (1, 2)
|
|
) # Last split should have 1 row
|
|
|
|
# Check that original metadata is preserved
|
|
for split_table in split_tables:
|
|
self.assertEqual(split_table.metadata, "metadata")
|
|
self.assertEqual(split_table.page_num, "1")
|
|
self.assertEqual(split_table.header, "Long Table")
|
|
|
|
def test_table_split_no_split_needed(self):
|
|
"""Test splitting a table that does not need splitting. It should return the original table."""
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["A", "B"], ["1", "2"], ["3", "4"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
split_tables = table.split(row_limit=10) # Row limit larger than table size
|
|
self.assertEqual(len(split_tables), 1) # Should return the original table
|
|
self.assertEqual(
|
|
split_tables[0].table.shape, (2, 2)
|
|
) # Should be the same shape as original
|
|
self.assertEqual(split_tables[0], table) # Should be the same object
|
|
|
|
def test_table_to_list_format(self):
|
|
"""Test conversion back to list format."""
|
|
table_data = [["A", "B"], ["1", "3"], ["2", "4"]]
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
list_format = table.to_list_format()
|
|
|
|
self.assertEqual(list_format, table_data)
|
|
self.assertIsInstance(list_format, list)
|
|
self.assertIsInstance(list_format[0], list) # Should be a list of lists
|
|
|
|
def test_table_info_new_format(self):
|
|
"""Test getting table info from new list-of-rows format."""
|
|
table_dict = get_table_info(self.sample_text_dict_new)
|
|
|
|
self.assertIn("1", table_dict)
|
|
self.assertEqual(len(table_dict["1"]), 1)
|
|
self.assertEqual(table_dict["1"][0].header, "Table 1")
|
|
self.assertEqual(table_dict["1"][0].metadata, "Some metadata")
|
|
self.assertEqual(
|
|
table_dict["1"][0].table.shape, (2, 2)
|
|
) # 2 rows, 2 columns in first table
|
|
self.assertEqual(table_dict["2"][0].header, "") # Second table has no header
|
|
self.assertEqual(
|
|
table_dict["2"][0].table.shape, (3, 2)
|
|
) # 3 rows, 2 columns in second table (empty header means all data rows)
|
|
|
|
def test_table_info_old_format(self):
|
|
"""Test getting table info from old dict format."""
|
|
with self.assertLogs(level="INFO") as log: # Expect a deprecation warning
|
|
table_dict = get_table_info(self.sample_text_dict_old)
|
|
|
|
self.assertIn("1", table_dict)
|
|
self.assertEqual(len(table_dict["1"]), 1)
|
|
self.assertEqual(table_dict["1"][0].header, "Table 1")
|
|
self.assertEqual(table_dict["1"][0].metadata, "Some metadata")
|
|
self.assertEqual(table_dict["1"][0].table.shape, (2, 2))
|
|
|
|
def test_get_table_info_multiple_tables(self):
|
|
"""Test getting table info from a page with multiple tables."""
|
|
table_dict = get_table_info(self.multi_table_text)
|
|
|
|
self.assertIn("1", table_dict)
|
|
self.assertEqual(len(table_dict["1"]), 2)
|
|
self.assertEqual(table_dict["1"][0].header, "First Table")
|
|
|
|
def test_recreate_page(self):
|
|
"""Test recreating page text from tables."""
|
|
table_dict = get_table_info(self.sample_text_dict_new)
|
|
page_tables = table_dict["1"]
|
|
recreated_text = recreate_page(page_tables)
|
|
|
|
self.assertIn("Table 1", recreated_text)
|
|
self.assertIn("-------Table Start--------", recreated_text)
|
|
self.assertIn("-------Table End--------", recreated_text)
|
|
self.assertIn("Some metadata", recreated_text)
|
|
|
|
def test_split_tables_by_row_count(self):
|
|
"""Test splitting tables that exceed row limit and adding .0 suffix."""
|
|
# Create a long table that will need splitting
|
|
long_table_data = [["Name", "Value"]] + [
|
|
[f"Row{i}", str(i)] for i in range(1, 8)
|
|
] # 7 data rows
|
|
text_dict = {
|
|
"5": f"Table metadata\n-------Table Start--------\nLong Table\n{long_table_data}\n-------Table End--------\nPost-table text"
|
|
}
|
|
|
|
result = split_tables_by_rows(text_dict, row_limit=3)
|
|
|
|
# Should create multiple pages with .0, .1, .2 suffixes
|
|
expected_pages = ["5.0", "5.1", "5.2"]
|
|
for page in expected_pages:
|
|
self.assertIn(page, result)
|
|
|
|
# Original page should not exist
|
|
self.assertNotIn("5", result)
|
|
|
|
# Last page should have post-table text
|
|
self.assertIn("Post-table text", result["5.2"])
|
|
|
|
def test_split_tables_by_row_count_no_split_needed(self):
|
|
"""Test that small tables get .0 suffix but aren't split."""
|
|
text_dict = {
|
|
"3": "Metadata\n-------Table Start--------\nSmall Table\n[['A', 'B'], ['1', '2']]\n-------Table End--------"
|
|
}
|
|
|
|
result = split_tables_by_rows(text_dict, row_limit=10)
|
|
|
|
# Should have .0 suffix but no additional splits
|
|
self.assertIn("3.0", result)
|
|
self.assertNotIn("3", result)
|
|
self.assertNotIn("3.1", result)
|
|
|
|
def test_split_tables_by_row_count_no_tables(self):
|
|
"""Test that pages without tables keep original page numbers."""
|
|
text_dict = {
|
|
"1": "Just some regular text without any tables.",
|
|
"2": "More text, no tables here either.",
|
|
}
|
|
|
|
result = split_tables_by_rows(text_dict, row_limit=5)
|
|
|
|
# Pages without tables should keep original numbers
|
|
self.assertIn("1", result)
|
|
self.assertIn("2", result)
|
|
self.assertNotIn("1.0", result)
|
|
self.assertNotIn("2.0", result)
|
|
|
|
def test_combine_continuous_tables_no_continuation(self):
|
|
"""Test that non-continuation tables are not combined."""
|
|
text_dict = {
|
|
"1": "Metadata 1\n-------Table Start--------\nTable 1\n[['A', 'B'], ['1', '2']]\n-------Table End--------",
|
|
"2": "Metadata 2\n-------Table Start--------\nTable 2\n[['C', 'D'], ['3', '4']]\n-------Table End--------",
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict)
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
|
|
# Both pages should remain (no continuation tables to combine)
|
|
self.assertIn("1", result)
|
|
self.assertIn("2", result)
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_table_combine_post_table_text_accumulation(self):
|
|
"""Test that post-table text is properly accumulated when combining tables."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="meta",
|
|
header="Main Table",
|
|
table_data=[["A", "B"], ["1", "2"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
post_table_text="First footnote",
|
|
)
|
|
|
|
table2 = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=[["3", "4"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
post_table_text="Second footnote",
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Both footnotes should be preserved
|
|
self.assertIn("First footnote", table1.post_table_text)
|
|
self.assertIn("Second footnote", table1.post_table_text)
|
|
|
|
def test_table_combine_column_mismatch_preserves_post_table_text(self):
|
|
"""Test that post-table text is preserved when combining tables with column mismatches."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Main Table",
|
|
table_data=[["A", "B"], ["1", "2"], ["x", "y"]], # Header + 2 data rows
|
|
start_index=0,
|
|
end_index=100,
|
|
post_table_text="Main table footnote",
|
|
)
|
|
|
|
table2 = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=[["3", "4", "5"]], # Just 1 data row (3 columns - mismatch!)
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
post_table_text="Continuation footnote",
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Verify both post-table texts are preserved despite column mismatch
|
|
self.assertIn("Main table footnote", table1.post_table_text)
|
|
self.assertIn("Continuation footnote", table1.post_table_text)
|
|
|
|
# Verify the table combination worked correctly
|
|
self.assertEqual(
|
|
table1.table.shape[0], 3
|
|
) # 2 original + 1 continuation = 3 total rows
|
|
self.assertEqual(table1.table.shape[1], 2) # 2 columns (from main table)
|
|
|
|
# Verify data integrity
|
|
self.assertEqual(table1.table.iloc[0, 0], "1") # Original data row 1
|
|
self.assertEqual(table1.table.iloc[1, 0], "x") # Original data row 2
|
|
self.assertEqual(table1.table.iloc[2, 0], "3") # Continuation data (truncated)
|
|
|
|
def test_table_combine_column_mismatch_padding_preserves_post_table_text(self):
|
|
"""Test post-table text preservation when padding columns."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Main Table",
|
|
table_data=[
|
|
["A", "B", "C"],
|
|
["1", "2", "3"],
|
|
["x", "y", "z"],
|
|
], # Header + 2 data rows
|
|
start_index=0,
|
|
end_index=100,
|
|
post_table_text="Main footnote",
|
|
)
|
|
|
|
table2 = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=[["4"]], # Just 1 data row (1 column - needs padding)
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
post_table_text="Padded footnote",
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Both footnotes should be preserved even with column padding
|
|
self.assertIn("Main footnote", table1.post_table_text)
|
|
self.assertIn("Padded footnote", table1.post_table_text)
|
|
|
|
# Verify padding worked
|
|
self.assertEqual(
|
|
table1.table.shape, (3, 3)
|
|
) # 2 original + 1 continuation = 3 rows, 3 columns
|
|
self.assertEqual(table1.table.iloc[0, 0], "1") # Original data row 1
|
|
self.assertEqual(table1.table.iloc[1, 0], "x") # Original data row 2
|
|
self.assertEqual(table1.table.iloc[2, 0], "4") # Continuation data
|
|
self.assertEqual(table1.table.iloc[2, 1], "") # Padded
|
|
self.assertEqual(table1.table.iloc[2, 2], "") # Padded
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|