How can I make my function to read a csv file and return a list in Python?-CodePudding

I have a function that takes 2 inputs. I want it to read a csv file and return a one dimensional list that contains data from every row in a successive manner. I want my function to list every vote counts for every nominee in a list.

Codes of my function is: `

import csv
import pandas as pd
def retreiveData(filename,nominees):
    file=open(filename)
    csv_reader=csv.reader(file)
    result=[]

    return result
nominees=["Obama","Romney","Johnson","Stein","Others"]
retreiveData("ElectionUSA2012.csv",nominees)

` There is some data in "ElectionUSA2012.csv".It contains vote counts for each nominee.

CodePudding user response：

why dont you use pd.read_csv() ? And convert dataframe to list:

    df = pd.read_csv(filename)
    result = df.to_numpy().tolist() #if you want list or numpy array

CodePudding user response：

How about this one?

import pandas as pd

def retrieve_data(csv_file_path, nominees):
    df = pd.read_csv(csv_file_path)
    # https://stackoverflow.com/questions/11346283/renaming-column-names-in-pandas
    counts = df.sum(axis = 0, skipna = True).values.tolist()
    # https://www.geeksforgeeks.org/python-pandas-dataframe-sum/
    counts = {k:v for k, v in zip(nominees, counts)}
    
    return counts

csv_file_path = 'csv.csv'
counts = retrieve_data(csv_file_path, nominees=["Obama", "Romney", "Johnson", "Stein", "Others"])
print(counts)

for a csv file like this:

A   B   C   D    E
1   2   3   4    6
4   4   2   3    6
3   3   3   10   6

It would print a dictionary like this:

{'Obama': 8, 'Romney': 9, 'Johnson': 8, 'Stein': 17, "Others": 18}