jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/io-utils-docs (pull request #713)

Browse Source 
io_utils documentation

* add docstring to read_local function and handle file read errors gracefully by returning None on failure

* update read_local function to support .xlsx files and improve error handling

* add docstring to read_s3 function to clarify its purpose and return value

* refactor read_s3 and read_s3_csv functions to improve error handling and add type hints

* improve error handling in read_s3_csv function for CSV file parsing

* enhance read_input function with error handling and detailed docstring

* Add TODOs for read_input_csv and filter_already_processed functions; add type hints and docstring for remove_txt_extension

* enhance write_local function documentation; make docstrings more consistent when there is a "notes" section

* enhance write_local and write_s3 functions; improve documentation, add error output type handling, and ensure consistent CSV writing options

* refactor test_io_utils.py; reorganize imports and enhance read_input_local test with additional mock for file existence


Approved-by: Katon Minhas
This commit is contained in:
Alex Galarce
2025-09-26 20:47:55 +00:00
parent 814d3e0636
commit 0562b90646
2 changed files with 186 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fieldExtraction/src/utils/io_utils.py
+168 -33
View File
@@ -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 "<BATCH_ID>-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 "<BATCH_ID>-RESULTS.csv"
in a timestamped directory under CONSOLIDATED_OUTPUT_DIRECTORY
- "individual" output_type saves the file as "<filename>_-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 "<BATCH_ID>-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
fieldExtraction/tests/test_io_utils.py
+18 -5
View File
@@ -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()