diff --git a/fieldExtraction/src/utils/io_utils.py b/fieldExtraction/src/utils/io_utils.py index f70a934..b00ba7a 100644 --- a/fieldExtraction/src/utils/io_utils.py +++ b/fieldExtraction/src/utils/io_utils.py @@ -4,14 +4,22 @@ import tempfile from io import StringIO import pandas as pd +import src.tracking.tracking_utils as tracking_utils from botocore.exceptions import ClientError from pyxlsb import open_workbook - -import src.tracking.tracking_utils as tracking_utils from src import config -def read_local(file_path): # io_utils.py +def read_local(file_path) -> str | pd.DataFrame | None: + """Reads a local (.txt, .csv, or .xlsx) file and returns its contents. + + Args: + file_path (str): The path to the local file. + + Returns: + str or pd.DataFrame: The contents of the file, either as a string (for text files) + or as a pandas DataFrame (for CSV/XLSX files). + """ # Check if the file is a text file if os.path.isfile(file_path) and file_path.endswith(".txt"): try: @@ -30,6 +38,7 @@ def read_local(file_path): # io_utils.py logging.error( f"Failed to decode {file_path} with UTF-8 and cp1252 encodings." ) + return None if os.path.isfile(file_path) and file_path.endswith(".csv"): try: @@ -37,6 +46,7 @@ def read_local(file_path): # io_utils.py return df except Exception: logging.error("Failed to read .csv file.") + return None if os.path.isfile(file_path) and file_path.endswith(".xlsx"): try: @@ -44,6 +54,11 @@ def read_local(file_path): # io_utils.py return df except Exception: logging.error("Failed to read .xlsx file.") + return None + + # Unsupported file type + logging.error(f"Unsupported file type for {file_path}.") + return None def list_s3_files(prefix=config.S3_PREFIX, bucket=config.S3_BUCKET): # io_utils.py @@ -79,7 +94,17 @@ def read_s3_pdf(pdf_path): return None -def read_s3(): +def read_s3() -> dict: + """Reads all .txt files from a specific S3 bucket and prefix. + + Returns: + dict: A dictionary containing the contents of the files, keyed by filename. + + Notes: + - Doesn't take input arguments; instead, it uses configurations from the config + module for S3 client, bucket, and prefix. + - Only reads files with a .txt extension. + """ s3_client = config.S3_CLIENT bucket = config.S3_BUCKET prefix = config.S3_PREFIX @@ -102,7 +127,9 @@ def read_s3(): if not obj["Key"].endswith("/") ] - if response.get("IsTruncated"): + if response.get( + "IsTruncated", False + ): # default to False to avoid infinite loop continuation_token = response.get("NextContinuationToken") else: break @@ -119,41 +146,93 @@ def read_s3(): context = contents.decode("utf-8") path, filename = os.path.split(contract) files[filename] = context - except: + except Exception as e: + logging.warning(f"Failed to read {contract}: {e}") pass return files -def read_s3_csv(csv_path): # io_utils.py +def read_s3_csv(csv_path) -> pd.DataFrame | None: + """Reads a CSV file from an S3 bucket and returns it as a DataFrame. + + Args: + csv_path (str): The S3 URL of the CSV file (e.g., s3://bucket-name/path/to/file.csv). + + Returns: + pd.DataFrame | None: The contents of the CSV file as a pandas DataFrame, + or None if an error occurs. + """ s3_client = config.S3_CLIENT - s3_parts = csv_path[5:].split("/", 1) + + # Remove 's3://' prefix and split into bucket and key + if not csv_path.startswith("s3://"): + logging.error(f"Invalid S3 path: {csv_path}") + return None + + s3_parts = csv_path[5:].split("/", 1) # Remove 's3://' prefix with magic number 5 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"]: + response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix) + for obj in response.get( + "Contents", [] + ): # If "Contents" key is missing, default to empty list 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 + csv_object = s3_client.get_object(Bucket=bucket, Key=key) + try: + df = pd.read_csv( + StringIO(csv_object["Body"].read().decode("utf-8")) + ) + return df + except Exception as e: + logging.error(f"Error parsing CSV file {key}: {e}") + continue except ClientError as e: logging.error(f"Error listing S3 objects: {e}") + return None + + # If no CSV files were found + logging.warning(f"No CSV files found at {csv_path}") + return None -def read_input(local_path=config.LOCAL_PATH): # io_utils.py +def read_input(local_path=config.LOCAL_PATH) -> dict: + """Reads input files from the specified local path or S3 bucket. + + Args: + local_path (str, optional): The local path to read files from. Defaults to config.LOCAL_PATH. + CAUTION: If config.READ_MODE is set to "s3", this parameter is ignored. + + Returns: + dict: A dictionary containing the contents of the input files. + + Notes: + - This function relies on the global configuration setting config.READ_MODE + - If config.READ_MODE is "local", it reads all files from the specified local path. + - If config.READ_MODE is "s3", it reads files from the configured S3 bucket and prefix. + """ 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 + try: + for file in os.listdir(local_path): + full_path = os.path.join(local_path, file) + if os.path.isfile(full_path): # Only process files and not directories + file_text = read_local(full_path) + files[file] = file_text + return files + except (FileNotFoundError, PermissionError, NotADirectoryError) as e: + logging.error(f"Error accessing local path {local_path}: {e}") + return {} elif config.READ_MODE == "s3": return read_s3() + else: + logging.error(f"Invalid READ_MODE: {config.READ_MODE}") + return {} -def read_input_csv(): # io_utils.py +def read_input_csv(): # TODO: This function is unreferenced outside of unit tests. Consider removal. 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 @@ -164,7 +243,9 @@ def read_input_csv(): # io_utils.py return ac_df, b_df -def filter_already_processed(input_dict, input_dict_b): # io_utils.py +def filter_already_processed( + input_dict, input_dict_b +): # TODO: this function is unreferenced outside of unit tests. Consider removal. """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 @@ -221,13 +302,23 @@ def filter_already_processed(input_dict, input_dict_b): # io_utils.py return input_dict_ac, input_dict_b -def remove_txt_extension(filename): # io_utils.py +def remove_txt_extension(filename: str) -> str: + """Remove the .txt extension from a filename. + + Args: + filename (str): The filename from which to remove the extension. + + Returns: + str: The filename, with .txt extension removed if it was present. + """ if isinstance(filename, str) and filename.lower().endswith(".txt"): return filename[:-4] return filename -def read_xlsb(path): # io_utils.py +def read_xlsb( + path, +): # TODO: this function is unreferenced outside of unit tests. Consider removal. with open_workbook(path) as wb: with wb.get_sheet(1) as sheet: rows = [] @@ -239,44 +330,83 @@ def read_xlsb(path): # io_utils.py return df -def write_local(df, filename, run_timestamp, output_type): +def write_local( + df: pd.DataFrame, filename: str, run_timestamp: str, output_type: str +) -> None: """ Saves a DataFrame as a CSV file in a local directory. Args: df (pd.DataFrame): The DataFrame to save. + filename (str): The base filename to use when saving the file. run_timestamp (str): The timestamp used to organize output files. + output_type (str): Type of output, either "final" or "individual". - The function creates a directory based on the configured output path - and the provided timestamp, then saves the DataFrame as a CSV file - named "-RESULTS.csv" within it. + Returns: + None; this function is called for its side effect of writing a file to local. + + Notes: + - Creates output directories if they do not exist + - Uses different naming conventions based on the output_type + - "final" output_type saves the file as "-RESULTS.csv" + in a timestamped directory under CONSOLIDATED_OUTPUT_DIRECTORY + - "individual" output_type saves the file as "_-RESULTS.csv" + in the OUTPUT_DIRECTORY """ if output_type == "final": output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp) os.makedirs(output_dir, exist_ok=True) - df.to_csv(os.path.join(output_dir, f"{config.BATCH_ID}-RESULTS.csv"), quoting=1) + df.to_csv( + os.path.join(output_dir, f"{config.BATCH_ID}-RESULTS.csv"), + index=False, + quoting=1, + ) + return None elif output_type == "individual": base_filename = os.path.splitext(filename)[0].strip() os.makedirs(config.OUTPUT_DIRECTORY, exist_ok=True) df.to_csv( os.path.join(config.OUTPUT_DIRECTORY, f"{base_filename}_-RESULTS.csv"), index=False, + quoting=1, ) + return None + elif output_type == "error": + output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp) + os.makedirs(output_dir, exist_ok=True) + df.to_csv( + os.path.join(output_dir, f"{config.BATCH_ID}-ERRORS.csv"), + index=False, + quoting=1, + ) + return None + else: + logging.error( + f"Unknown output_type: {output_type}; output_type must be 'final', 'individual', or 'error'" + ) + return None -def write_s3(df, filename, run_timestamp, output_type): +def write_s3( + df: pd.DataFrame, filename: str, run_timestamp: str, output_type: str +) -> None: """ Uploads a DataFrame to an S3 bucket as a CSV file. Args: df (pd.DataFrame): The DataFrame to be written to S3. + filename (str): The base filename to use when saving the file. run_timestamp (str): A timestamp representing the current run. + output_type (str): Type of output, either "final", "individual", or "error". The function generates an S3 key based on the configured BATCH_ID and uploads the DataFrame as a CSV file to the configured S3 bucket. The CSV file is saved in the "consolidated" folder with the naming convention "-RESULTS.csv". + Returns: + None; this function is called for its side effect of uploading a file to S3. + Configurations used: - config.BATCH_ID: Identifier for the batch process. - config.S3_CLIENT: The AWS S3 client for performing operations. @@ -302,10 +432,15 @@ def write_s3(df, filename, run_timestamp, output_type): config.S3_CLIENT.put_object( Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer ) + if output_type == "final": + logging.info(f"Saved batch output file: {output_path}") + elif output_type == "individual": + logging.info(f"Saved individual output file: {output_path}") + elif output_type == "error": + logging.info(f"Saved error output file: {output_path}") + else: + logging.info(f"Saved unknown output type file: {output_path}") + return None except ClientError as e: logging.error(f"Error uploading to S3: {e}") - - if output_type == "final": - logging.info(f"Saved batch output file: {output_path}") - elif output_type == "individual": - logging.info(f"Saved individual output file: {output_path}") + return None diff --git a/fieldExtraction/tests/test_io_utils.py b/fieldExtraction/tests/test_io_utils.py index 415ffe1..eb17c66 100644 --- a/fieldExtraction/tests/test_io_utils.py +++ b/fieldExtraction/tests/test_io_utils.py @@ -2,11 +2,17 @@ from io import StringIO import pandas as pd import pytest - -from src.utils.io_utils import (filter_already_processed, list_s3_files, - read_input, read_input_csv, read_local, - read_s3, read_s3_csv, read_xlsb, - remove_txt_extension) +from src.utils.io_utils import ( + filter_already_processed, + list_s3_files, + read_input, + read_input_csv, + read_local, + read_s3, + read_s3_csv, + read_xlsb, + remove_txt_extension, +) class TestIOUtils: @@ -118,6 +124,7 @@ class TestIOUtils: # Tests for read_input def test_read_input_local(self, mock_txt_data, mocker): mocker.patch("os.listdir", return_value=["file1.txt"]) + mocker.patch("os.path.isfile", return_value=True) mocker.patch("src.utils.io_utils.read_local", return_value=mock_txt_data) mocker.patch("src.config.READ_MODE", "local") result = read_input() @@ -296,3 +303,9 @@ class TestIOUtils: # Verify that the mock was called correctly mock_open_workbook.assert_called_once_with("dummy_path.xlsb") mock_workbook.get_sheet.assert_called_once() + # Assert that the result matches the expected DataFrame + pd.testing.assert_frame_equal(result, expected_df) + + # Verify that the mock was called correctly + mock_open_workbook.assert_called_once_with("dummy_path.xlsb") + mock_workbook.get_sheet.assert_called_once()