Alex Galarce
0562b90646
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
447 lines
16 KiB
Python
447 lines
16 KiB
Python
import logging
|
|
import os
|
|
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
|
|
from src import config
|
|
|
|
|
|
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:
|
|
# 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
|
|
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:
|
|
df = pd.read_csv(file_path)
|
|
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:
|
|
df = pd.read_excel(file_path)
|
|
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
|
|
"""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_pdf(pdf_path):
|
|
"""
|
|
Downloads a PDF file from an S3 bucket and returns a temporary file path.
|
|
|
|
Args:
|
|
pdf_path (str): The S3 URL of the PDF file (e.g., s3://bucket-name/path/to/file.pdf).
|
|
|
|
Returns:
|
|
str: The local temporary file path of the downloaded PDF.
|
|
"""
|
|
|
|
s3_client = config.S3_CLIENT
|
|
s3_parts = pdf_path[5:].split("/", 1)
|
|
bucket = s3_parts[0]
|
|
prefix = s3_parts[1] if len(s3_parts) > 1 else ""
|
|
|
|
try:
|
|
response = s3_client.get_object(Bucket=bucket, Key=prefix)
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
|
|
temp_pdf.write(response["Body"].read())
|
|
temp_pdf_path = temp_pdf.name
|
|
return temp_pdf_path
|
|
except ClientError as e:
|
|
logging.error(f"Error fetching PDF from S3: {e}")
|
|
return None
|
|
|
|
|
|
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
|
|
|
|
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", False
|
|
): # default to False to avoid infinite loop
|
|
continuation_token = response.get("NextContinuationToken")
|
|
else:
|
|
break
|
|
|
|
contract_list = sorted(file_list)
|
|
files = {}
|
|
|
|
for contract in contract_list:
|
|
logging.info(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 Exception as e:
|
|
logging.warning(f"Failed to read {contract}: {e}")
|
|
pass
|
|
|
|
return files
|
|
|
|
|
|
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
|
|
|
|
# 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:
|
|
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"):
|
|
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) -> 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 = {}
|
|
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(): # 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
|
|
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 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
|
|
already_processed_ac, already_processed_b = tracking_utils.get_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
|
|
}
|
|
|
|
logging.info(
|
|
f"Filtering based on S3 master tracking for batch {config.BATCH_ID}:"
|
|
)
|
|
logging.info(
|
|
f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining"
|
|
)
|
|
logging.info(
|
|
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
|
|
}
|
|
|
|
logging.info(f"Filtering based on local output directory:")
|
|
logging.info(
|
|
f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining"
|
|
)
|
|
logging.info(
|
|
f"B: Found {len(already_processed_b)} already processed, {len(input_dict_b)} remaining"
|
|
)
|
|
|
|
return input_dict_ac, input_dict_b
|
|
|
|
|
|
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,
|
|
): # 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 = []
|
|
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
|
|
|
|
|
|
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".
|
|
|
|
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"),
|
|
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: 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.
|
|
- config.S3_OUTPUT_BUCKET: The target S3 bucket for storing the file.
|
|
|
|
Prints:
|
|
Confirmation message with the S3 key of the uploaded file.
|
|
"""
|
|
if output_type == "final":
|
|
output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-RESULTS.csv"
|
|
elif output_type == "individual":
|
|
output_path = (
|
|
f"{config.BATCH_ID}/{run_timestamp}/individual/{filename}-RESULTS.csv"
|
|
)
|
|
elif output_type == "error":
|
|
output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-ERRORS.csv"
|
|
else:
|
|
# Fallback path for unknown output types
|
|
output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-unknown-{output_type}.csv"
|
|
|
|
csv_buffer = df.to_csv(index=False).encode()
|
|
try:
|
|
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}")
|
|
return None
|