Terrence Rozay Terrence Rozay - 23 days ago 6
Python Question

change function to return entire column

The function returns a single state and population but I need to to return all the populations in a column but skip over the first 6 rows

def findpop(state=None):
f = open(getMediaPath("population_state_reduced (2).csv"), "rt")
index = 1
for line in f:
if index > 3:
parts = line.split(',')
if state is None:
return [(parts[4], int(parts[5]))]
else:
for line in f:
if parts[4] == state.capitalize():
return int(parts[5])
index += 1

print findpop()

Answer

If I understood you correct this is exactly what you want:

import csv

def findpop(state=None):
    res = []
    with open(getMediaPath("population_state_reduced (2).csv")) as f:
        reader = csv.reader(f)
        for i, line in enumerate(reader):
            if i > 5: # skip first 6 rows
                if state is None:
                    res.append((line[4], int(line[5])))
                else:
                    if line[4] == state.capitalize():
                        return int(line[5])
    return res