86 lines
3 KiB
Python
Executable file
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()
|
|
|