132 lines
3.7 KiB
Python
132 lines
3.7 KiB
Python
import sys
|
|
import os
|
|
import csv
|
|
import argparse
|
|
|
|
"""
|
|
|
|
Splits a CSV file into multiple files based on command line arguments.
|
|
|
|
Arguments:
|
|
|
|
`-h`: help file of usage of the script
|
|
`-i`: input file name
|
|
`-o`: output file name
|
|
`-r`: row limit to split
|
|
|
|
Default settings:
|
|
|
|
`output_path` is the current directory
|
|
headers are displayed on each split file
|
|
the default delimeter is a comma
|
|
|
|
Example usage:
|
|
|
|
```
|
|
# split csv by every 100 rows
|
|
>> python csv_split.py -i input.csv -o output -r 100
|
|
```
|
|
|
|
"""
|
|
|
|
|
|
def get_arguments():
|
|
"""Grab user supplied arguments using the argparse library."""
|
|
|
|
# Use arparse to get command line arguments
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-i", "--input_file", required=True,
|
|
help="csv input file (with extension)", type=str)
|
|
parser.add_argument("-o", "--output_file", required=True,
|
|
help="csv output file (without extension)", type=str)
|
|
parser.add_argument("-r", "--row_limit", required=True,
|
|
help="row limit to split csv at", type=int)
|
|
args = parser.parse_args()
|
|
|
|
# Check if the input_file exits
|
|
is_valid_file(parser, args.input_file)
|
|
|
|
# Check if the input_file is valid
|
|
is_valid_csv(parser, args.input_file, args.row_limit)
|
|
|
|
return args.input_file, args.output_file, args.row_limit
|
|
|
|
|
|
def is_valid_file(parser, file_name):
|
|
"""Ensure that the input_file exists."""
|
|
if not os.path.exists(file_name):
|
|
parser.error("The file '{}' does not exist!".format(file_name))
|
|
sys.exit(1)
|
|
|
|
|
|
def is_valid_csv(parser, file_name, row_limit):
|
|
"""
|
|
Ensure that the # of rows in the input_file
|
|
is greater than the row_limit.
|
|
"""
|
|
row_count = 0
|
|
for row in csv.reader(open(file_name)):
|
|
row_count += 1
|
|
# Note: You could also use a generator expression
|
|
# and the sum() function to count the rows:
|
|
# row_count = sum(1 for row in csv.reader(open(file_name)))
|
|
if row_limit > row_count:
|
|
parser.error(
|
|
"The 'row_count' of '{}' is > the number of rows in '{}'!"
|
|
.format(row_limit, file_name)
|
|
)
|
|
sys.exit(1)
|
|
|
|
|
|
def parse_file(arguments):
|
|
"""
|
|
Splits the CSV into multiple files or chunks based on the row_limit.
|
|
Then create new CSV files.
|
|
"""
|
|
input_file = arguments[0]
|
|
output_file = arguments[1]
|
|
row_limit = arguments[2]
|
|
output_path = '.' # Current directory
|
|
|
|
# Read CSV, split into list of lists
|
|
with open(input_file, 'r') as input_csv:
|
|
datareader = csv.reader(input_csv)
|
|
all_rows = []
|
|
for row in datareader:
|
|
all_rows.append(row)
|
|
|
|
# Remove header
|
|
header = all_rows.pop(0)
|
|
|
|
# Split list of list into chunks
|
|
current_chunk = 1
|
|
for i in range(0, len(all_rows), row_limit): # Loop through list
|
|
chunk = all_rows[i:i + row_limit] # Create single chunk
|
|
|
|
current_output = os.path.join( # Create new output file
|
|
output_path,
|
|
"{}-{}.csv".format(output_file, current_chunk)
|
|
)
|
|
|
|
# Add header
|
|
chunk.insert(0, header)
|
|
|
|
# Write chunk to output file
|
|
with open(current_output, 'w') as output_csv:
|
|
writer = csv.writer(output_csv)
|
|
writer = writer.writerows(chunk)
|
|
|
|
# Output info
|
|
print ""
|
|
print "Chunk # {}:".format(current_chunk)
|
|
print "Filepath: {}".format(current_output)
|
|
print "# of rows: {}".format(len(chunk))
|
|
|
|
# Create new chunk
|
|
current_chunk += 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
arguments = get_arguments()
|
|
parse_file(arguments)
|