iframe-proxy | Sunbelt Computer Software

executable file
131 lines (100 loc) · 3.69 KB
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)
Sunbelt Computer Software

PL/B Language Development and Support

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

FilesExpand file tree

12_csv_split.py

Latest commit

History

12_csv_split.py

File metadata and controls
