如何在python中获取csv数据

debugcn 发表于 Dev

那敏

我已经编写了plot函数和主体部分，但是无法找出获取数据csv的必要代码。所以基本上我的代码是：

import random
import matplotlib.pyplot as plt 

def get_data_csv(fn, colx, coly, sep=','):


    return x_data, y_data, x_label, y_label

def plot_graphs(x_data, y_data, x_label, y_label, title, color=None):
    plt.figure()
    #if there are more than one graph to plot in the same figure, same x-axis
    if all(type(arg) == list for arg in [x_data, y_data, y_label]):
        lx = len(x_data)
        if lx > 1 and all(len(arg) == lx for arg in [y_data, y_label]):
            for i in range(lx):
                color = list(random.random() for c in range(3))
                plt.plot(x_data[i], y_data[i], color = color, label = y_label[i])
            plt.xlabel(x_label)
            plt.title(title)
            plt.legend()
        else:
            print('Something get wrong, review your arguments')
    # if there is only one graph to plot
    else:
        plt.plot(x_data, y_data, color=color)
        plt.ylabel(y_label)
        plt.xlabel(x_label)
        plt.title(title)
        plt.xlim([min(x_data), max(x_data)])
    plt.grid('all')
    plt.show()
    
if __name__ == "__main__":
    user_fn = input("Enter CSV filename(s), separated by ',': ")
    user_colx = input("Enter respective x-column(s), separated by ',': ")
    user_coly = input("Enter respective y-column(s), separated by ',': ")
    user_title = input("Enter title: ")
    fn = user_fn.split(',')
    colx = list(int(c) for c in user_colx.split(','))
    coly = list(int(c) for c in user_coly.split(','))
    if len(fn) == 1:
        fn = fn[0]
        x_data, y_data, x_label, y_label = get_data_csv(fn, colx[0], coly[0])
    else:
        x_data, y_data, y_label = [], [], []
        for i in range(len(fn)):
            auxx, auxy, x_label, auxylbl = get_data_csv(fn[i], colx[i], coly[i])
            x_data.append(auxx)
            y_data.append(auxy)
            y_label.append(auxylbl)
    plot_graphs(x_data, y_data, x_label, y_label, user_title, color='black')

如您所见，我无法编写该def get_data_csv部分。我认为我应该使用readline()从文件中获取标头行的方法，然后，我可以使用任何方式获取其他信息来举例说明代码应如何工作。

Enter CSV filename(s), separated by ',': sinusoid.csv
Enter respective x-column(s), separated by ',': 0
Enter respective y-column(s), separated by ',': 1
Enter title: Sine

这应该有效。我需要帮助

那敏

我自己找到了问题的答案，所以我需要的代码是：

def get_data_csv(fn, colx, coly, sep=','):
    x_data = []
    y_data = []
    with open(fn) as f:
        lines = f.readlines()
        i = 0
        for line in lines:
            line = line[0:-2]
            line = line.split(',')
            if i == 0:
                x_label = line[colx]
                y_label = line[coly]
            else:
                x_value = float(line[colx])
                x_data.append(x_value)
                y_value = float(line[coly])
                y_data.append(y_value)
            i+=1    
    return x_data, y_data, x_label, y_label

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。