数据绘图

import xlrd
import os
import csv
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import numpy as np


class MTSPLOT(object):
    def __init__(self):
        self.dl_5gdata = []
        self.colorflag = []
        self.color = []

    def get_mtsfilename(self):
        filelist = os.listdir()
        for e in filelist:
            if '.csv' in e:
                return e
        raise Exception

    def get_mtsdata(self):
        with open(self.get_mtsfilename()) as f:
            f_csv = csv.reader(f)
            print(type(f_csv))
            print(dir(f_csv))
            print(next(f_csv)[11])
            print(next(f_csv))
            for e in f_csv:
                self.dl_5gdata.append(float(e[11]))
                self.colorflag.append(e[2]+e[3])
            print("ok")

    def plot_data(self):
        ax = plt.axes()
        y = self.dl_5gdata
        print(y)
        print(max(y))
        #y = [1,4,9,16,25]
        x = range(len(y))
        setcolor = list(set(self.colorflag))
        linecolor = ['red','blue','black','green']
        for i in self.colorflag:
                self.color.append(linecolor[setcolor.index(i)])

        
        points = np.array([x,y]).T.reshape(-1,1,2)
        points = np.concatenate([points[:-1], points[1:]], axis=1)
        lc = LineCollection(points,color=self.color)
        ax.set_xlim(min(x), max(x))
        ax.set_ylim(min(y), max(y))
        ax.add_collection(lc)
        plt.show()





if __name__ == '__main__':
    mts = MTSPLOT()
    print(mts.get_mtsfilename())
    mts.get_mtsdata()
    mts.plot_data()