Search code examples
pythonjython

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()

Solution

  • 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