cs50p/problems/pset6/scourgify/scourgify.py

import csv
import os
import sys


def main():
    try:
        input_file_name, output_file_name = sys.argv[1:3]
    except IndexError:
        print("Too few command-line arguments")
        sys.exit(1)
    if len(sys.argv) > 3:
        print("Too many command-line arguments")
        sys.exit(1)
    if not os.path.isfile(input_file_name):  # Or try opening and catch FileNotFoundError
        print(f"Could not read {input_file_name}")
        sys.exit(1)
    if not input_file_name.endswith('.csv'):
        print(f"Not a CSV file")
        sys.exit(1)
    write_names_to_csv_file(read_names_from_csv_file(input_file_name), output_file_name)


def read_names_from_csv_file(input_file_name):
    with open(input_file_name, 'r') as input_file:
        reader = csv.DictReader(input_file)
        for row in reader:
            yield row


def write_names_to_csv_file(input_data, output_file_name):
    with open(output_file_name, 'w') as output_file:
        writer = csv.DictWriter(output_file, fieldnames=['first', 'last', 'house'])
        writer.writeheader()
        for row in input_data:
            row['last'], row['first'] = row['name'].split(', ')
            del row['name']
            writer.writerow(row)


if __name__ == "__main__":
    main()