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

84 lines
2.9 KiB
Python
Executable file

#!/usr/bin/env python
# Script Name: fill-in-missing-data
import os
import csv
from collections import defaultdict
from pyfiglet import Figlet
from halo import Halo
# Define the paths based on the project details
PROJECT_ROOT = "/home/ld/mgk-scrapes"
STAGE_6_DIRECTORY = os.path.join(PROJECT_ROOT, 'current-data', '.data', 'stage-6')
INPUT_CSV = os.path.join(STAGE_6_DIRECTORY, '12-re-ordered-columns.csv')
OUTPUT_CSV = os.path.join(STAGE_6_DIRECTORY, '13-filled-in-missing-data.csv')
def fill_in_missing_data():
"""Fill in missing data for each row of match group by fetching data from other rows in the group."""
columns_to_merge = ['YiB', 'GBP Business Category', 'GBP Matching Service']
uuid_to_rows = defaultdict(list)
# Read the input CSV and group rows by Business UUID
with open(INPUT_CSV, 'r') as infile:
reader = csv.DictReader(infile)
headers = reader.fieldnames
if 'Business UUID' not in headers:
print("Error: 'Business UUID' column is missing.")
return
rows = list(reader)
print(f"Total rows read from input file: {len(rows)}")
for row in rows:
business_uuid = row['Business UUID'].strip()
if business_uuid: # Ignore empty UUIDs
uuid_to_rows[business_uuid].append(row)
print(f"Total groups by Business UUID: {len(uuid_to_rows)}")
# Fill in missing data within each group
for uuid, grouped_rows in uuid_to_rows.items():
# Collect data from all rows in the group
merged_data = {column: '' for column in columns_to_merge}
for row in grouped_rows:
for column in columns_to_merge:
if row[column].strip():
merged_data[column] = row[column]
# Fill missing data in each row
for row in grouped_rows:
for column in columns_to_merge:
if not row[column].strip():
row[column] = merged_data[column]
# Write to the output CSV
with open(OUTPUT_CSV, 'w', newline='') as outfile:
writer = csv.DictWriter(outfile, fieldnames=headers)
writer.writeheader()
for row in rows:
writer.writerow(row)
print(f"Total rows written to output file: {len(rows)}")
def main():
figlet = Figlet(font='slant')
script_name = "fill-in-missing-data".replace("-", " ").title()
print(figlet.renderText(script_name))
# Check if the output file already exists and delete it if it does
if os.path.exists(OUTPUT_CSV):
os.remove(OUTPUT_CSV)
print(f"Existing output file {OUTPUT_CSV} found and deleted.")
spinner = Halo(text='Filling in missing data by Business UUID', spinner='dots')
spinner.start()
fill_in_missing_data()
spinner.succeed('Missing data filled successfully.')
print(f"Processed input CSV: {INPUT_CSV}")
print(f"Output CSV: {OUTPUT_CSV}")
if __name__ == '__main__':
main()