Adding rows to CSV sorted by ID in Python without reading whole file into memory

user_10654888 0
Viewed 283

I have a CSV file where the first column is an ID like so:

5,a
4,b
2,c
1,d

The rows must always be sorted from the biggest to smallest ID. I have a list of rows that I want to add, without reading the whole original CSV in memory, so I can't just append the rows and sort after. Here is the code I came up with:

import csv


def main():
    rows_to_add = [[7, "NEW1"], [6, "NEW2"], [3, "NEW3"], [-2, "NEW4"]]

    with open("in.csv", "r") as in_file, open("out.csv", "w") as out_file:
        reader = csv.reader(in_file)
        writer = csv.writer(out_file)

        for new_row in rows_to_add:
            for source_row in reader:
                if new_row[0] >= int(source_row[0]):
                    writer.writerow(new_row)
                    writer.writerow(source_row)
                    break

                writer.writerow(source_row)
            else:
                # If source reader already reached end of file
                writer.writerow(new_row)

        for remaining_line in in_file:
            out_file.write(remaining_line) 

    with open("out.csv", "r") as out_file:
        print(out_file.read())


if __name__ == "__main__":
    main()

Result:

7,NEW1
5,a
6,NEW2
4,b
3,NEW3
2,c
1,d
-2,NEW4

This doesn't work correctly if there are two consecutive IDs, 6,NEW2 should be just after 7,NEW1 and I can't figure out the right way to do it.

csv python
1 Answers

You have to use a running pointer for at least one of the lists. In this case, as you can't read the entire CSV, the running pointer can be used for the other list

Following code should work

import csv


def main():
    rows_to_add = [[7, "NEW1"], [6, "NEW2"], [6, "NEW2"], [3, "NEW3"], [-2, "NEW4"]]

    with open("in.csv", "r") as in_file, open("out.csv", "w") as out_file:
        reader = csv.reader(in_file)
        writer = csv.writer(out_file)
        idx = 0

        # For each line in file
        for source_row in reader:
            # Print all rows from rows_to_add that are larger
            while rows_to_add[idx][0] > int(source_row[0]):
                writer.writerow(rows_to_add[idx])
                idx += 1

            # Before printing current line from file
            writer.writerow(source_row)

        # Print remaining rows in rows_to_add
        for row in rows_to_add[idx:]:
            writer.writerow(row)


    with open("out.csv", "r") as out_file:
        print(out_file.read())


if __name__ == "__main__":
    main()

sample output for your in.csv

7,NEW1
6,NEW2
6,NEW2
5,a
4,b
3,NEW3
2,c
1,d
-2,NEW4
user_18380760
Related