CSV Data Processor
This utility reads a CSV file, filters rows based on a specific criteria, selects specific columns, and saves the result to a new file. It demonstrates how to use csv.DictReader and csv.DictWriter for robust and memory-efficient CSV handling.
Modules Used:
The Code
Save this as csv_tool.py.
import csv
import argparse
import sys
def process_csv(input_file, output_file, filter_col=None, filter_val=None, columns=None):
try:
# Open input and output files
with open(input_file, mode='r', newline='', encoding='utf-8') as infile, \
open(output_file, mode='w', newline='', encoding='utf-8') as outfile:
reader = csv.DictReader(infile)
if not reader.fieldnames:
print("Error: CSV file is empty or has no header.")
return
# Determine output columns (use all if none specified)
out_fields = columns if columns else reader.fieldnames
# Validate that requested columns exist in input
for col in out_fields:
if col not in reader.fieldnames:
print(f"Error: Column '{col}' not found in input CSV.")
return
# Initialize writer
writer = csv.DictWriter(outfile, fieldnames=out_fields, extrasaction='ignore')
writer.writeheader()
count = 0
for row in reader:
# 1. Filter Logic
if filter_col:
if filter_col not in row:
print(f"Error: Filter column '{filter_col}' not found.")
return
# Skip row if it doesn't match the filter value
if row[filter_col] != filter_val:
continue
# 2. Write Row
# DictWriter with extrasaction='ignore' will automatically
# ignore keys in 'row' that aren't in 'out_fields'
writer.writerow(row)
count += 1
print(f"Success! Processed {count} rows into '{output_file}'.")
except FileNotFoundError:
print(f"Error: File '{input_file}' not found.")
except Exception as e:
print(f"An error occurred: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="CSV Data Processor")
parser.add_argument("input", help="Input CSV file path")
parser.add_argument("output", help="Output CSV file path")
parser.add_argument("--filter-col", help="Column name to filter by")
parser.add_argument("--filter-val", help="Value to match in the filter column")
parser.add_argument("--columns", nargs="+", help="Specific columns to keep (space separated)")
args = parser.parse_args()
process_csv(args.input, args.output, args.filter_col, args.filter_val, args.columns)Usage
# Filter rows where 'Status' is 'Active' and keep only 'Name' and 'Email' columns
python csv_tool.py users.csv active_users.csv --filter-col Status --filter-val Active --columns Name Email