doczyai-pipelines/textract-pipeline/src/lambda/prompt-orchestrator/table_funcs.py

import re
import ast


def clean_tables(text):
    def replace_colon(text):
        # Define a function to use in re.sub to check the context of the match
        def replacer(match):
            # Extract the character after ':' to check if it's '['
            following_text = match.group(1)
            if following_text.strip().startswith("["):
                return match.group(
                    0
                )  # Return the original match (':' and whatever follows)
            else:
                return "" + match.group(1)  # Replace ':' with ';' and return the rest

        # Use a regular expression to find ':' and the text that follows
        pattern = r":(\s*.)"
        replaced_text = re.sub(pattern, replacer, text)
        return replaced_text

    def remove_unquoted_brackets(text):
        # Define a function to check the content between brackets
        def replacer(match):
            content = match.group(1)
            if "'" not in content and (len(content) > 0):
                return content  # Return just the content without brackets
            else:
                return match.group(
                    0
                )  # Return the original match if it contains a quote

        # Use a regular expression to find brackets and the text within
        pattern = r"\[([^\[\]]*?)\]"
        cleaned_text = re.sub(pattern, replacer, text)
        return cleaned_text

    # Clean tables
    text = text.strip("{}")  # Remove brackets
    text = replace_colon(text)
    text = remove_unquoted_brackets(text)
    return text


def convert_to_dict(table_text):
    """
    Converts a formatted table text into a dictionary where each key maps to a list of values.

    This function processes a string representation of a table, extracting key-value pairs and converting them into a dictionary.
    Each key maps to a list of values parsed from the input string. The function also returns the number of elements in each list.

    Parameters:
        table_text (str): The input string representing the table, formatted with keys and values.

    Returns:
        tuple: A dictionary with keys mapping to lists of values, and an integer representing the number of elements in each list.
    """
    table_text = clean_tables(table_text)
    # print(table_text, '\n')
    table_text_list = table_text.split("]")
    table_text_list = [item for item in table_text_list if len(item) > 0]

    final_dict = {}
    num_elements = 0
    for key_value_text in table_text_list:
        key = key_value_text.split(":")[0].strip(", '")
        try:
            value = ":".join(key_value_text.split(":")[1:]).strip() + "]"
            # print(value, '\n')
            if len(value) > 1:
                value_list = ast.literal_eval(value)
                final_dict[key] = value_list
                num_elements = len(value_list)
        except:
            value = ":".join(key_value_text.split(":")[1:]).strip() + "']"
            # print(value, '\n')
            if len(value) > 1:
                value_list = ast.literal_eval(value)
                final_dict[key] = value_list
                num_elements = len(value_list)

    return final_dict, num_elements


def format_table(table_json, table_size):
    """
    Formats a dictionary of table data into a string representation of the table.

    This function takes a dictionary where each key maps to a list of values and converts it into a formatted string.
    The keys are used as headers, and each row is constructed by iterating through the lists up to the specified table size.
    The resulting string separates headers and values with ' : ' and includes newline characters to delineate rows.

    Parameters:
        table_json (dict): The dictionary containing the table data, with keys as column headers and lists as column values.
        table_size (int): The number of rows in the table.

    Returns:
        str: The formatted string representation of the table.
    """
    table_text = ""
    keys = [key for key in table_json.keys()]
    table_text += " : ".join(keys) + "\n"
    for i in range(table_size):
        for key in keys:
            if table_json[key][i]:
                # table_text += key + ': ' + table_json[key][i] + ', '
                table_text += table_json[key][i] + " : "
        table_text += r"\n"
    return table_text


def align_and_format_tables(text_dict):
    """
    Aligns and formats tables within the text of each page in a dictionary.

    This function searches for table text within the input text of each page, marked by 'Table Start' and 'Table End' tags.
    It extracts, converts, and formats these tables, then replaces the original table text with the formatted version while
    maintaining the alignment with surrounding text. The cleaned and formatted text for each page is then stored in a new dictionary.

    Parameters:
        text_dict (dict): A dictionary where keys are page numbers and values are the text content of those pages, potentially containing tables.

    Returns:
        dict: A new dictionary with aligned and formatted tables in the text for each page.
    """
    aligned_text_dict = {}
    for key, text in text_dict.items():
        aligned_text = text
        if "Table Start" in text:
            table_texts = re.findall(
                r"-------Table Start--------(.*?)-------Table End--------",
                text,
                re.DOTALL,
            )
            for table_text in table_texts:
                # print(table_text, '\n')

                try:
                    # Extract pretable text
                    pretable = table_text.split("{")[0].strip()
                    # Extract and format table text
                    # table_only = "{" + table_text.split('{', 1)[1].rsplit('}', 1)[0].replace("'", '"') + "}"
                    table_only = (
                        "{" + table_text.split("{", 1)[1].rsplit("}", 1)[0] + "}"
                    )

                    table, table_size = convert_to_dict(table_only)
                    table_formatted = format_table(table, table_size)
                except:
                    table_formatted = table_text

                # Align
                if text.count(pretable) == 2:  # One table
                    # Remove original table
                    aligned_text = text.replace(table_text, "")
                    # Align new table
                    aligned_text = aligned_text.replace(
                        pretable, pretable + table_formatted
                    )

                else:
                    aligned_text = aligned_text.replace(
                        table_text, pretable + " " + table_formatted
                    )
            aligned_text_dict[key] = aligned_text.replace(
                "-------Table Start--------", ""
            ).replace("-------Table End--------", "")
        else:
            aligned_text_dict[key] = text

    return aligned_text_dict