76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
# MIT License
|
|
# Copyright (c) 2024 Jeff Moe
|
|
""" Read CSV and extract selected columns and write to new CVS"""
|
|
|
|
import csv
|
|
import argparse
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Extract CSV Columns, Output CSV")
|
|
|
|
parser.add_argument("headers_file", help="Headers File", type=str)
|
|
parser.add_argument("input_csv", help="Input CSV File", type=str)
|
|
parser.add_argument("output_csv", help="Output CSV File", type=str)
|
|
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def read_good_headers(filename):
|
|
"""Reads and returns the list of 'good' headers from a given file."""
|
|
with open(filename, "r") as file:
|
|
return [line.strip() for line in file.readlines()]
|
|
|
|
|
|
def filter_csv(input_csv, output_csv, good_headers):
|
|
"""Filters an input CSV based on the provided good headers and writes to output CSV."""
|
|
|
|
# Read the original CSV
|
|
with open(input_csv, mode="r", newline="", encoding="utf-8") as infile:
|
|
reader = csv.DictReader(infile)
|
|
|
|
# Get only the required fieldnames ('good' headers)
|
|
filtered_fieldnames = [
|
|
field for field in reader.fieldnames if field in good_headers
|
|
]
|
|
|
|
# Write to output CSV
|
|
with open(output_csv, mode="w", newline="", encoding="utf-8") as outfile:
|
|
writer = csv.DictWriter(outfile, fieldnames=filtered_fieldnames)
|
|
|
|
# Write the header line (column names) first
|
|
writer.writeheader()
|
|
|
|
for row in reader:
|
|
filtered_row = {
|
|
key: value
|
|
for key, value in row.items()
|
|
if key in filtered_fieldnames
|
|
}
|
|
writer.writerow(filtered_row)
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
headers_file = args.headers_file
|
|
input_csv = args.input_csv
|
|
output_csv = args.output_csv
|
|
|
|
# input_file = "tmp/mycontacts.csv"
|
|
# good_headers_file = "tmp/good_headers.txt"
|
|
# output_file = "tmp/final.csv"
|
|
|
|
# Step 1: Read the list of good headers
|
|
good_headers = read_good_headers(headers_file)
|
|
|
|
# Step 2: Filter the CSV based on these headers and write to a new file
|
|
filter_csv(input_csv, output_csv, good_headers)
|
|
|
|
print(f"Filtered CSV has been written to {output_csv}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|