instant-data-scraper-utils/bin/stage-6/prepare-stage-6.py
2024-07-16 18:08:36 -04:00

86 lines
3 KiB
Python
Executable file

#!/usr/bin/env python
# Script Name: prepare-stage-6
import os
import shutil
import csv
from halo import Halo
from tqdm import tqdm
from pyfiglet import Figlet
# 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_5_DIRECTORY = os.path.join(DATA_DIRECTORY, "stage-5")
STAGE_6_DIRECTORY = os.path.join(DATA_DIRECTORY, "stage-6")
MERGED_CSV = os.path.join(STAGE_6_DIRECTORY, "01-first-merger.csv")
HEADERS = [
"Business Name",
"GBP Review Rating",
"GBP Review Count",
"GBP Business Category",
"YiB",
"GBP Matching Service",
"GBP Business Phone",
"GBP Business Website",
"GBP Location Municipality",
"GBP Location State",
"Root Domain",
"MD5 for Scrape" # Added the new column header
]
def create_stage_6_directory():
"""Create the stage-6 directory, replacing it if it already exists."""
if os.path.exists(STAGE_6_DIRECTORY):
print("Existing Stage 6 found, deleting and replacing...")
shutil.rmtree(STAGE_6_DIRECTORY)
os.makedirs(STAGE_6_DIRECTORY)
print("Stage 6 directory created.")
def merge_csv_files():
"""Merge all CSV files from stage-5 into a single CSV file in stage-6."""
with open(MERGED_CSV, 'w', newline='') as outfile:
writer = csv.DictWriter(outfile, fieldnames=HEADERS)
writer.writeheader()
for state_dir in os.listdir(STAGE_5_DIRECTORY):
state_path = os.path.join(STAGE_5_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)
with open(file_path, 'r') as infile:
reader = csv.DictReader(infile)
for row in reader:
data = {header: row.get(header, '').strip() for header in HEADERS}
writer.writerow(data)
spinner.succeed(f'Finished processing {state_dir}')
def main():
figlet = Figlet(font='slant')
script_name = "prepare-stage-6".replace("-", " ").title()
print(figlet.renderText(script_name))
print("Preparing Stage 6...")
create_stage_6_directory()
print("Merging CSV files to Stage 6...")
spinner = Halo(text='Merging data', spinner='dots')
spinner.start()
merge_csv_files()
spinner.succeed("Data merged.")
print("Stage 6 preparation complete.")
if __name__ == "__main__":
main()