instant-data-scraper-utils/bin/stage-2/sanitize-review-count.py
2024-07-16 18:08:36 -04:00

85 lines
2.8 KiB
Python
Executable file

#!/usr/bin/env python
# Script Name: sanitize-review-count
import os
import csv
import re
import sys
from pyfiglet import Figlet
from halo import Halo
# Define the paths based on the project details
PROJECT_ROOT = "/home/ld/mgk-scrapes"
CURRENT_DATASET = os.path.join(PROJECT_ROOT, "current-data")
DATA_DIRECTORY = os.path.join(CURRENT_DATASET, ".data")
STAGE_2_DIRECTORY = os.path.join(DATA_DIRECTORY, "stage-2")
def sanitize_review_count_data(file_path):
"""Sanitize review count data in the given CSV file."""
with open(file_path, 'r') as file:
reader = csv.reader(file)
rows = list(reader)
if not rows:
return
headers = rows[0]
data_rows = rows[1:]
review_count_pattern = re.compile(r'^-\d{1,3}(?:,\d{3})*$')
review_count_index = None
for i, header in enumerate(headers):
if any(review_count_pattern.match(row[i]) for row in data_rows):
review_count_index = i
break
if review_count_index is None:
return
modified_rows = []
for row in rows:
if review_count_pattern.match(row[review_count_index]):
row[review_count_index] = f"({row[review_count_index][1:]})"
modified_rows.append(row)
with open(file_path, 'w', newline='') as file:
writer = csv.writer(file)
writer.writerows(modified_rows)
def process_csv_files(stage_directory):
"""Process all CSV files in the given stage directory to sanitize review count data."""
for state_dir in os.listdir(stage_directory):
state_path = os.path.join(stage_directory, state_dir)
if os.path.isdir(state_path):
spinner = Halo(text=f'Processing {state_dir}', spinner='dots')
spinner.start()
for county_dir in os.listdir(state_path):
county_path = os.path.join(state_path, county_dir)
if os.path.isdir(county_path):
csv_files = [file for file in os.listdir(county_path) if file.endswith('.csv')]
for file in csv_files:
file_path = os.path.join(county_path, file)
sanitize_review_count_data(file_path)
spinner.succeed(f'Finished processing {state_dir}')
def main():
figlet = Figlet(font='slant')
script_name = "sanitize-review-count".replace("-", " ").title()
print(figlet.renderText(script_name))
if not os.path.exists(STAGE_2_DIRECTORY):
print(f"Error: The directory {STAGE_2_DIRECTORY} does not exist.")
sys.exit(1)
print(f"Sanitizing review counts in {STAGE_2_DIRECTORY}...")
process_csv_files(STAGE_2_DIRECTORY)
final_spinner = Halo(spinner='dots', color='green')
final_spinner.start()
final_spinner.succeed("Review counts sanitized.")
if __name__ == "__main__":
main()