doczyai-pipelines/fieldExtraction/src/utils.py

import os
import re
import pandas as pd
import shutil
import numpy as np
import tracking 
import config
from botocore.exceptions import ClientError
from io import StringIO
from pyxlsb import open_workbook


class InvalidDateException(Exception):
    pass


def read_local(file_path): # io_utils.py
    # Check if the file is a text file
    if os.path.isfile(file_path) and file_path.endswith(".txt"):
        try:
            # First attempt to open the file with UTF-8 encoding
            with open(file_path, "r", encoding="utf-8") as file:
                file_contents = file.read()
                return file_contents
        except UnicodeDecodeError:
            # If UTF-8 fails, try reading the file with ANSI encoding
            try:
                with open(file_path, "r", encoding="cp1252") as file:
                    file_contents = file.read()
                    return file_contents
            except UnicodeDecodeError:
                # If ANSI also fails, log an error message or handle it accordingly
                print(f"Failed to decode {file_path} with UTF-8 and cp1252 encodings.")
    
    if os.path.isfile(file_path) and file_path.endswith(".csv"):
        try:
            df = pd.read_csv(file_path)
            return df
        except Exception:
            print('Failed to read .csv file.')

    if os.path.isfile(file_path) and file_path.endswith(".xlsx"):
        try:
            df = pd.read_excel(file_path)
            return df
        except Exception:
            print('Failed to read .xlsx file.')


def list_s3_files(prefix=config.S3_PREFIX, bucket=config.S3_BUCKET): # io_utils.py
    """List all files in an S3 folder."""
    response = config.S3_CLIENT.list_objects_v2(Bucket=bucket, Prefix=prefix)
    return [item["Key"] for item in response.get("Contents", [])]


def read_s3():
    s3_client = config.S3_CLIENT
    bucket = config.S3_BUCKET
    prefix = config.S3_PREFIX

    file_list = []
    continuation_token = None

    while True:

        if continuation_token:
            response = s3_client.list_objects_v2(
                Bucket=bucket, Prefix=prefix, ContinuationToken=continuation_token
            )
        else:
            response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)

        file_list += [
            obj["Key"]
            for obj in response.get("Contents", [])
            if not obj["Key"].endswith("/")
        ]

        if response.get("IsTruncated"):
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    contract_list = sorted(file_list)
    files = {}

    for contract in contract_list:
        print(f"Reading {contract}")
        try:
            data = s3_client.get_object(Bucket=bucket, Key=contract)
            if contract.endswith('.txt'):
                contents = data["Body"].read()
                context = contents.decode("utf-8")
                path, filename = os.path.split(contract)
                files[filename] = context
        except:
            pass

    return files

def read_s3_csv(csv_path): # io_utils.py
    s3_client = config.S3_CLIENT
    s3_parts = csv_path[5:].split("/", 1)
    bucket = s3_parts[0]
    prefix = s3_parts[1] if len(s3_parts) > 1 else ""
    try:
        for obj in s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)["Contents"]:
            key = obj["Key"]
            if key.endswith(".csv"):
                response = s3_client.get_object(Bucket=bucket, Key=key)
                df = pd.read_csv(StringIO(response["Body"].read().decode("utf-8")))
                return df
    except ClientError as e:
        print(f"Error listing S3 objects: {e}")

def read_input(local_path=config.LOCAL_PATH): # io_utils.py
    if config.READ_MODE == "local":
        files = {}
        for file in os.listdir(local_path):
            full_path = os.path.join(local_path, file)
            file_text = read_local(full_path)
            files[file] = file_text
        return files
    elif config.READ_MODE == "s3":
        return read_s3()
    
def read_input_csv(): # io_utils.py
    if config.DF_READ_MODE == "local":
        ac_df = read_local(config.AC_DF) if config.AC_DF else None
        b_df = read_local(config.B_DF) if config.B_DF else None
    elif config.DF_READ_MODE == "s3":
        ac_df =  read_s3_csv(config.AC_DF) if config.AC_DF else None
        b_df = read_s3_csv(config.B_DF) if config.B_DF else None

    return ac_df, b_df

def calculate_progress_stats(completed_ac, total_ac, completed_b, total_b): # tracking.py
    """Calculate and return progress statistics"""
    ac_progress = min((completed_ac / total_ac * 100), 100.0) if total_ac > 0 else 0
    b_progress = min((completed_b / total_b * 100), 100.0) if total_b > 0 else 0
    return ac_progress, b_progress

def contains_reimbursement(text, page="1"): # string_funcs.py
    if isinstance(text, dict):
        return page.isdigit() and (
            "%" in text[page]
            or "$" in text[page]
            or "percent " in text[page]
            or "compensation schedule" in text[page].lower()
            or "reimbursement schedule" in text[page].lower()
        )
    elif isinstance(text, str):
        return (
            "%" in text
            or "$" in text
            or "percent " in text
            or "compensation schedule" in text.lower()
            or "reimbursement schedule" in text.lower()
        )
    else:
        print("contains_reimbursement - Invalid data type")


def filter_already_processed(input_dict, input_dict_b): # io_utils.py
    """Filter out already processed files based on either local files or S3 master tracking"""
    if config.WRITE_TO_S3:
        # Get processed files from S3 master tracking for this specific batch
        already_processed_ac, already_processed_b = tracking.get_already_processed_files(config.BATCH_ID)
        
        # Keep files that still need either AC or B processing
        input_dict_ac = {
            k: v for k, v in input_dict.items() 
            if k not in already_processed_ac
        }
        
        input_dict_b = {
            k: v for k, v in input_dict_b.items() 
            if k not in already_processed_b
        }
        
        print(f"Filtering based on S3 master tracking for batch {config.BATCH_ID}:")
        print(f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining")
        print(f"B: Found {len(already_processed_b)} already processed, {len(input_dict_b)} remaining")
        
    else:
        # Original local directory checking logic
        already_processed_ac, already_processed_b = [], []
        for folder_name in os.listdir(config.OUTPUT_DIRECTORY):
            folder_path = os.path.join(config.OUTPUT_DIRECTORY, folder_name)
            if os.path.isdir(folder_path):
                if config.AC_RESULTS_NAME in os.listdir(folder_path):
                    already_processed_ac.append(folder_name + ".txt")
                if config.B_RESULTS_NAME in os.listdir(folder_path):
                    already_processed_b.append(folder_name + ".txt")
        
        input_dict_ac = {
            k: v for k, v in input_dict.items()
            if k not in already_processed_ac
        }
        
        input_dict_b = {
            k: v for k, v in input_dict_b.items()
            if k not in already_processed_b
        }
        
        print(f"Filtering based on local output directory:")
        print(f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining")
        print(f"B: Found {len(already_processed_b)} already processed, {len(input_dict_b)} remaining")

    return input_dict_ac, input_dict_b




def is_empty(value): # string_funcs.py
    if pd.isna(value):
        return True
    else:
        empty_values = [None, "", "N/A", "NA", "null", "none", "NaN", np.nan, "nan"]
        return value in empty_values


def find_regex_matches( # regex_funcs.py
    pattern: str, text_dict: dict
) -> tuple[list, list]:  # TODO: Write few unit tests
    """Helper function to find matches and list of pages of a given regex pattern in text_dict.

    Args:
        pattern(str): A particular regex match that with a pattern.
        text_dict (dict): Dictionary keyed by string page-number and valued by actual textract output page

    Returns:
        matches(list): List of match(s)
        page_list(list): List of pages where the match was found

    """
    matches = []
    page_list = []
    for page, text in text_dict.items():
        found = re.findall(pattern, text)
        if found:
            matches.extend(found) # make it append
            page_list.append(page)

    return matches, page_list



def add_hyphen_if_needed(input_str): # postprocessing_funcs.py
    """Checks for specific patterns in the input string and modifies it by adding 
    a '0' at the start and a hyphen at the 3rd index, or just adds a hyphen at 
    the 3rd index if it's a 9-digit number.

    If the input string is exactly 8 characters long, it adds a '0' at the start 
    and inserts a hyphen at the 3rd position. If the string matches the pattern 
    of a 9-digit number (exactly 9 digits), it inserts a hyphen at the 3rd index.
    If the input string matches the pattern 'x-xxxxxxx', it adds '0' at the start.

    All conditions will return a string in the format 'xx-xxxxxxx'.

    Args:
    - input_str (str): The input string (IRS).

    Returns:
    - str: The modified string in the format 'xx-xxxxxxx'.
    
    Example:
    - '1-2345678' → '01-2345678'
    - '12345678' → '01-2345678'
    - '123456789' → '12-3456789'
    """
    # Check if the string matches the pattern 'x-xxxxxxx' (1 character, hyphen, 7 characters)
    if re.fullmatch(r"\d-\d{7}", input_str):
        # Add '0' at the start
        input_str = '0' + input_str
    
    # Check if the string is exactly 'xxxxxxxx' (8 characters)
    elif bool(re.fullmatch(r"\b\d{8}\b", input_str)):
        # Add '0' at the start and insert a hyphen at the 3rd index
        input_str = '0' + input_str
        input_str = input_str[:2] + '-' + input_str[2:]
    
    # If the string matches the pattern of a 9-digit number, add a hyphen at the 3rd index
    elif bool(re.fullmatch(r"\b\d{9}\b", input_str)):
        input_str = input_str[:2] + '-' + input_str[2:]
    
    # Ensure the output is in the format 'xx-xxxxxxx'
    if not bool(re.fullmatch(r"\b\d{2}-\d{7}\b", input_str)):
        input_str = None

    return input_str

def convert_to_us_date_format(date_str: str) -> str: # postprocessing_funcs.py
    """
    Input date string is expected to be in the format YYYY-MM-DD
    
    Ensure the date is in MM/DD/YYYY format. If the month is greater than 12, swap the month and day.
        If the input is `N/A` or some variant (defined by is_empty), InvalidDateException will be raised.
    Args:
        date_str (str): The date string to be formatted. We expect it to come in as YYYY-MM-DD format (note the hyphens as well)

    Returns:
        str: The date string formatted as MM/DD/YYYY.

    Raises:
        InvalidDateException: If the date string is not in the expected format.
    """
    
    if not isinstance(date_str, str):
        raise TypeError(f"Date must be a string, got: {type(date_str)}")

    if is_empty(date_str) or re.fullmatch(r"\d{4}-\d{1,2}-\d{1,2}", date_str) is None:
        raise InvalidDateException(f"Invalid date format: {date_str}")

    # Split the date string into year, month, and day
    year, month, day = date_str.split("-")
    
    year_int = int(year)
    month_int = int(month)
    day_int = int(day)
    
    if year_int == 0 or month_int == 0 or day_int == 0:
        raise InvalidDateException(f"Invalid date format: {date_str}")
    
    if month_int > 12 and 0 < day_int <= 31:
        # Swap the month and day if the month is greater than 12
        month_int, day_int = day_int, month_int
    
    if 0 < month_int <= 12 and 0 < day_int <= 31:
        # Return the formatted date in MM/DD/YYYY format
        year, month, day = str(year_int).zfill(4), str(month_int).zfill(2), str(day_int).zfill(2)
        return "/".join([month, day, year])
    
    else:
        # Raise an exception if the month or day is out of range
        raise InvalidDateException(f"Invalid date format: {date_str}")


def remove_unnamed_columns(df): # postprocessing_funcs.py
    return df[[col for col in df.columns if 'Unnamed' not in col]]

def remove_txt_extension(filename): # io_utils.py
    if isinstance(filename, str) and filename.lower().endswith(".txt"):
        return filename[:-4]
    return filename

def read_xlsb(path): # io_utils.py
    with open_workbook(path) as wb:
        with wb.get_sheet(1) as sheet:
            rows = []
            for row in sheet.rows():
                rows.append([item.v for item in row])  # Extracting cell values
            df = pd.DataFrame(rows[1:], columns=rows[0])  # Converting to DataFrame, assuming first row is header
    return df