108 lines
3.8 KiB
Python
Executable file
108 lines
3.8 KiB
Python
Executable file
#!/usr/bin/env python
|
|
# Script Name: delete-empty-counties
|
|
|
|
import os
|
|
import sys
|
|
import shutil
|
|
import halo
|
|
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")
|
|
|
|
def get_stage_directory(stage):
|
|
"""Get the path for the specified stage directory."""
|
|
return os.path.join(DATA_DIRECTORY, stage)
|
|
|
|
def find_empty_counties(stage_directory):
|
|
"""Find all empty county directories in the given stage directory."""
|
|
empty_counties = []
|
|
for state_dir in os.listdir(stage_directory):
|
|
state_path = os.path.join(stage_directory, state_dir)
|
|
if os.path.isdir(state_path):
|
|
for county_dir in os.listdir(state_path):
|
|
county_path = os.path.join(state_path, county_dir)
|
|
if os.path.isdir(county_path):
|
|
# Check if the county directory contains any CSV files
|
|
csv_files = [f for f in os.listdir(county_path) if f.endswith('.csv')]
|
|
if not csv_files:
|
|
empty_counties.append(county_path)
|
|
return empty_counties
|
|
|
|
def find_empty_states(stage_directory):
|
|
"""Find all empty state directories in the given stage directory."""
|
|
empty_states = []
|
|
for state_dir in os.listdir(stage_directory):
|
|
state_path = os.path.join(stage_directory, state_dir)
|
|
if os.path.isdir(state_path):
|
|
county_dirs = [d for d in os.listdir(state_path) if os.path.isdir(os.path.join(state_path, d))]
|
|
if not county_dirs:
|
|
empty_states.append(state_path)
|
|
else:
|
|
# Check if all counties are empty
|
|
if all(os.path.isdir(os.path.join(state_path, county_dir)) and
|
|
not any(f.endswith('.csv') for f in os.listdir(os.path.join(state_path, county_dir)))
|
|
for county_dir in county_dirs):
|
|
empty_states.append(state_path)
|
|
return empty_states
|
|
|
|
def delete_directories(directories):
|
|
"""Delete the given directories."""
|
|
for directory in directories:
|
|
shutil.rmtree(directory)
|
|
|
|
def main():
|
|
figlet = Figlet(font='slant')
|
|
script_name = "delete-empty-counties".replace("-", " ").title()
|
|
print(figlet.renderText(script_name))
|
|
|
|
if len(sys.argv) != 2:
|
|
print("Usage: delete-empty-counties.py <stage-directory>")
|
|
sys.exit(1)
|
|
|
|
stage = sys.argv[1]
|
|
stage_directory = get_stage_directory(stage)
|
|
|
|
if not os.path.exists(stage_directory):
|
|
print(f"Error: The directory {stage_directory} does not exist.")
|
|
sys.exit(1)
|
|
|
|
print(f"Scanning for empty counties and states in {stage_directory}...")
|
|
spinner = halo.Halo(text='Scanning', spinner='dots')
|
|
spinner.start()
|
|
empty_counties = find_empty_counties(stage_directory)
|
|
empty_states = find_empty_states(stage_directory)
|
|
spinner.succeed("Scan complete.")
|
|
|
|
if empty_counties:
|
|
print("Empty counties found:")
|
|
print(", ".join(empty_counties))
|
|
|
|
print("Deleting empty counties...")
|
|
spinner.start()
|
|
delete_directories(empty_counties)
|
|
spinner.succeed("County deletion complete.")
|
|
|
|
print(f"Total county directories deleted: {len(empty_counties)}")
|
|
else:
|
|
print("No empty counties found.")
|
|
|
|
if empty_states:
|
|
print("Empty states found:")
|
|
print(", ".join(empty_states))
|
|
|
|
print("Deleting empty states...")
|
|
spinner.start()
|
|
delete_directories(empty_states)
|
|
spinner.succeed("State deletion complete.")
|
|
|
|
print(f"Total state directories deleted: {len(empty_states)}")
|
|
else:
|
|
print("No empty states found.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|