113 lines
4 KiB
Python
Executable file
113 lines
4 KiB
Python
Executable file
#!/usr/bin/env python
|
|
# Script Name: remove-empty-columns
|
|
|
|
import os
|
|
import csv
|
|
import sys
|
|
from pyfiglet import Figlet
|
|
import halo
|
|
|
|
def find_project_root(current_path):
|
|
"""Find the project root by locating the 'bin' directory."""
|
|
while current_path != os.path.dirname(current_path):
|
|
if os.path.basename(current_path) == 'bin':
|
|
return os.path.dirname(current_path)
|
|
current_path = os.path.dirname(current_path)
|
|
raise FileNotFoundError("Could not find 'bin' directory in the path hierarchy.")
|
|
|
|
# Define the paths based on the project details
|
|
script_path = os.path.abspath(__file__)
|
|
PROJECT_ROOT = find_project_root(script_path)
|
|
CURRENT_DATASET = os.path.join(PROJECT_ROOT, "current-data")
|
|
DATA_DIRECTORY = os.path.join(CURRENT_DATASET, ".data")
|
|
|
|
def print_help_message():
|
|
"""Prints the help message for script usage."""
|
|
print("Usage: remove-empty-columns.py <stage-#>")
|
|
print("Example: remove-empty-columns.py stage-1")
|
|
print("The provided stage directory must exist within the .data directory of the project root.")
|
|
|
|
def delete_empty_columns(file_path):
|
|
"""Deletes empty columns from a CSV file."""
|
|
with open(file_path, 'r') as file:
|
|
reader = csv.reader(file)
|
|
rows = list(reader)
|
|
|
|
if not rows:
|
|
return 0
|
|
|
|
headers = rows[0]
|
|
data_rows = rows[1:]
|
|
|
|
# Ensure each row has the same length as the headers
|
|
for row in data_rows:
|
|
while len(row) < len(headers):
|
|
row.append('')
|
|
|
|
# Identify non-empty columns
|
|
non_empty_columns = [i for i in range(len(headers)) if any(row[i].strip() for row in data_rows)]
|
|
empty_column_count = len(headers) - len(non_empty_columns)
|
|
|
|
# Create a new list to store the modified rows without empty columns
|
|
modified_rows = [[headers[i] for i in non_empty_columns]]
|
|
for row in data_rows:
|
|
modified_rows.append([row[i] for i in non_empty_columns])
|
|
|
|
# Write the modified rows back to the CSV file
|
|
with open(file_path, 'w', newline='') as file:
|
|
writer = csv.writer(file)
|
|
writer.writerows(modified_rows)
|
|
|
|
return empty_column_count
|
|
|
|
def remove_empty_columns(stage_directory):
|
|
"""Remove empty columns from all CSV files in the given stage directory."""
|
|
for state_dir in os.listdir(stage_directory):
|
|
state_path = os.path.join(stage_directory, state_dir)
|
|
if os.path.isdir(state_path):
|
|
print(f"Searching for empty columns to delete in {state_dir}...")
|
|
spinner = halo.Halo(spinner='dots', color='green')
|
|
spinner.start()
|
|
|
|
total_removed_columns = 0
|
|
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)
|
|
total_removed_columns += delete_empty_columns(file_path)
|
|
|
|
spinner.succeed(f"Removed columns from {state_dir}. Columns removed: {total_removed_columns}")
|
|
|
|
def main():
|
|
figlet = Figlet(font='slant')
|
|
script_name = "remove-empty-columns".replace("-", " ").title()
|
|
print(figlet.renderText(script_name))
|
|
|
|
if len(sys.argv) != 2:
|
|
print("Error: Incorrect number of arguments.")
|
|
print_help_message()
|
|
sys.exit(1)
|
|
|
|
stage_arg = sys.argv[1]
|
|
if not stage_arg.startswith('stage-'):
|
|
print("Error: Invalid stage directory format.")
|
|
print_help_message()
|
|
sys.exit(1)
|
|
|
|
stage_dir = os.path.join(DATA_DIRECTORY, stage_arg)
|
|
if not os.path.exists(stage_dir):
|
|
print(f"Error: The directory '{stage_dir}' does not exist.")
|
|
sys.exit(1)
|
|
|
|
print(f"Removing empty columns in {stage_arg}...")
|
|
remove_empty_columns(stage_dir)
|
|
|
|
final_spinner = halo.Halo(spinner='dots', color='green')
|
|
final_spinner.succeed("Finished removing empty columns.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|