117 lines
4.9 KiB
Python
Executable file
117 lines
4.9 KiB
Python
Executable file
#!/usr/bin/env python
|
|
# Script Name: md5-the-scrapes
|
|
|
|
import os
|
|
import hashlib
|
|
import pandas as pd
|
|
import pyfiglet
|
|
from halo import Halo
|
|
from tqdm import tqdm
|
|
from multiprocessing import Pool, cpu_count
|
|
|
|
# Function to calculate MD5 hash
|
|
def calculate_md5(row, exclude_columns):
|
|
md5_hash = hashlib.md5()
|
|
for column in row.index:
|
|
if column not in exclude_columns:
|
|
md5_hash.update(str(row[column]).encode('utf-8'))
|
|
return md5_hash.hexdigest()
|
|
|
|
# Function to determine if a column should be excluded based on patterns
|
|
def should_exclude(column_data):
|
|
# Ensure the column data is treated as strings
|
|
column_data = column_data.astype(str)
|
|
|
|
# Exclude column if it contains 5-star ratings (e.g., "4.1", "5.0")
|
|
if column_data.str.contains(r'^[1-5]\.\d$').any():
|
|
return True
|
|
# Exclude column if it contains review counts (e.g., "(10)", "(1,234)", "-10")
|
|
if column_data.str.contains(r'^\(\d{1,3}(?:,\d{3})*\)$').any() or column_data.str.contains(r'^-\d+$').any():
|
|
return True
|
|
# Exclude column if it contains years in business (e.g., "5+ years in business")
|
|
if column_data.str.contains(r'^\d+\+ years in business$').any():
|
|
return True
|
|
return False
|
|
|
|
# Function to check if a column contains terms from a given list
|
|
def contains_terms(column_data, terms):
|
|
column_data = column_data.astype(str)
|
|
match_count = column_data.apply(lambda x: any(term in x for term in terms)).sum()
|
|
return match_count / len(column_data) > 0.5 # Consider column matching if more than 50% cells match terms
|
|
|
|
# Function to process a single state directory
|
|
def process_state(state_dir_info):
|
|
state_dir, stage_1_directory, gbp_matching_services, gbp_business_categories = state_dir_info
|
|
state_path = os.path.join(stage_1_directory, state_dir)
|
|
spinner = Halo(text=f'Processing state: {state_dir}', spinner='dots')
|
|
spinner.start()
|
|
|
|
try:
|
|
for root, dirs, files in os.walk(state_path):
|
|
for file in files:
|
|
if file.endswith(".csv"):
|
|
file_path = os.path.join(root, file)
|
|
df = pd.read_csv(file_path)
|
|
|
|
# Identify columns to exclude
|
|
exclude_columns = []
|
|
|
|
# Identify GBP Business Categories column
|
|
gbp_business_col = None
|
|
for column in df.columns:
|
|
if contains_terms(df[column], gbp_business_categories):
|
|
gbp_business_col = column
|
|
break
|
|
|
|
# Identify GBP Matching Services column
|
|
for column in df.columns:
|
|
if column != gbp_business_col and contains_terms(df[column], gbp_matching_services):
|
|
exclude_columns.append(column)
|
|
break
|
|
|
|
# Identify other columns to exclude based on patterns
|
|
for column in df.columns:
|
|
if should_exclude(df[column]):
|
|
exclude_columns.append(column)
|
|
|
|
# Calculate MD5 for each row and add to new column "MD5 for Scrape"
|
|
df['MD5 for Scrape'] = df.apply(lambda row: calculate_md5(row, exclude_columns), axis=1)
|
|
|
|
# Save the modified CSV
|
|
df.to_csv(file_path, index=False)
|
|
|
|
spinner.succeed(f'Processing of state {state_dir} completed successfully.')
|
|
except Exception as e:
|
|
spinner.fail(f'Error processing state {state_dir}: {e}')
|
|
|
|
def main():
|
|
# Print script name using pyfiglet
|
|
script_name = "md5-the-scrapes".replace("-", " ").title()
|
|
print(pyfiglet.figlet_format(script_name))
|
|
|
|
# Define the directory structure
|
|
project_root = "/home/ld/mgk-scrapes"
|
|
stage_1_directory = os.path.join(project_root, "current-data", ".data", "stage-1")
|
|
gbp_matching_services_path = os.path.join(project_root, "current-data", ".data", "gbp-matching-services.txt")
|
|
gbp_business_categories_path = os.path.join(project_root, "current-data", ".data", "gbp-business-categories.txt")
|
|
|
|
# Load GBP Matching Services and Business Categories lists
|
|
with open(gbp_matching_services_path, 'r') as file:
|
|
gbp_matching_services = file.read().splitlines()
|
|
with open(gbp_business_categories_path, 'r') as file:
|
|
gbp_business_categories = file.read().splitlines()
|
|
|
|
# Get the list of state directories
|
|
state_dirs = [d for d in os.listdir(stage_1_directory) if os.path.isdir(os.path.join(stage_1_directory, d))]
|
|
|
|
# Prepare arguments for parallel processing
|
|
state_dir_info_list = [(state_dir, stage_1_directory, gbp_matching_services, gbp_business_categories) for state_dir in state_dirs]
|
|
|
|
# Process each state directory in parallel using multiprocessing
|
|
with Pool(cpu_count()) as pool:
|
|
pool.map(process_state, state_dir_info_list)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|