Comparing columns from two CSV files - follow-up

Question

This is the second part of a question, you can find the first part here:

Comparing columns from two CSV files

I've made some changes to the script and here is what it looks like now:

import csv, sys

def get_column(columns, name):
    count = 0
    for column in columns:
        if column != name:
            count += 1
        else:
            return(count)

def set_up_file(file, variable):
    columns = next(file)
    siren_pos = get_column(columns, 'SIREN')
    nic_pos = get_column(columns, 'NIC')
    variable_pos = get_column(columns, variable)
    return(siren_pos, nic_pos, variable_pos)

def test_variable(variable):
    with open('source.csv', 'r') as source:
        source_r = csv.reader(source, delimiter=';')
        sir_s, nic_s, comp_s = set_up_file(source_r, variable)
        line_s = next(source_r)
        with open('tested.csv', 'r') as tested:
            tested_r = csv.reader(tested, delimiter=';')
            sir_t, nic_t, comp_t = set_up_file(tested_r, variable)
            size = sum(1 for line in tested_r)
            tested.seek(0, 0)
            line_t = next(tested_r)
            line_t = next(tested_r)
            correct = 0
            try:
                while True:
                    if(line_s[sir_s] == line_t[sir_t] 
                    and line_s[nic_s] == line_t[nic_t]):
                        if(line_s[comp_s] == line_t[comp_t]):
                            correct += 1
                        line_t = next(tested_r)
                        line_s = next(source_r)
                    elif(int(line_s[sir_s]) > int(line_t[sir_t])):
                        line_t = next(tested_r)
                    elif(int(line_s[sir_s]) < int(line_t[sir_t])):
                        line_s = next(source_r)
                    else:
                        if(int(line_s[nic_s]) > int(line_t[nic_t])):
                            line_t = next(tested_r)
                        else:
                            line_s = next(source_r)
            except StopIteration:
                return(correct / size * 100)

def main():
    with open('tested.csv', 'r') as file:
        file_r = csv.reader(file, delimiter=';')
        columns = next(file_r)
        found = test_variable('SIREN')
        for column in columns:
            print(column, '%.2f' % (test_variable(column) / found * 100))

if(__name__ == "__main__"):
    main()

There are no real performance issue in the new version but I feel like it could still be improved greatly.

Also, would it be possible to reduce the size of the test_variable function? I thought about cutting it right before the try statement but I'll end up passing about 7 parameters which is not really a clean solution in my opinion.

Answer 1

• There is some duplication of code because you are handling two files.
• When you need some data converted to int it would be best to do soon as you've read it to avoid sprinkling other logic with int() calls.
• The two columns SIREN and NIC combined seem to form a sorting key to the file. You could simplify the if elif... part by performing the comparisons on (SIREN, NIC) tuples.

To address the above, I propose to organize the code like this:

def parse_file(file, variable):
    reader = csv.reader(file, delimiter=';')
    sir_s, nic_s, comp_s = set_up_file(reader, variable)
    for line in reader:
        key = int(line[sir_s]), int(line[nic_s])
        yield key, line[comp_s]

def test_variable(variable):
    with open('source.csv', 'r') as source, open('tested.csv', 'r') as tested:
        source_r = parse_file(source, variable)
        tested_r = parse_file(tested, variable)

        correct = 0
        try:
            line_s = next(source_r)
            line_t = next(tested_r)
            while True:
                key_s, comp_s = line_s
                key_t, comp_t = line_t
                if line_s == line_t:
                    correct += 1
                if key_s >= key_t:
                    line_t = next(tested_r)
                if key_s <= key_t:
                    line_s = next(source_r)

        except StopIteration:
            return correct

Note however that I've omitted the computation of size. This could be done by incrementing a variable after reading each line, but since lines are read at multiple places, and some may be left in the end if the other file ends first, it may be best to count the lines separately like you have done.

Answer 2

On this line:

return(siren_pos, nic_pos, variable_pos)

you do not need to surround the return values in parens. They will be returned as a tuple anyway.

It seems ill-advised to iterate through the entire file just to count the number of lines. I advise incrementing size as you go along, so that you only have to iterate once.

You have a bunch of common code that applies to both files. I recommend writing function that does the open(), the csv.reader(), and the set_up_file().