source: misc/tools/plotcsv/plotcsv.py @ 6477

Last change on this file since 6477 was 6477, checked in by rwilson, 15 years ago

Fixed bug whereby spaces in CSV header strings crashed the application.

File size: 20.6 KB
Line 
1#!/usr/bin/env python
2
3'''A module to draw a graph on the screen given a data file.
4
5The code is written such that it may be run as a program or imported
6and used as a module.
7'''
8
9import sys
10import os.path
11import types
12import csv
13import time
14import getopt
15try:
16    import cpickle
17except:
18    import pickle
19
20try:
21    import pylab
22except ImportError:
23    import tkinter_error
24    msg = ('Sorry, you have to install matplotlib.\n'
25           'Get it from either:\n'
26           r'   N:\georisk\downloads\matplotlib, or' '\n'
27           '   [http://matplotlib.sourceforge.net/].\n\n'
28           'You will also need to have numpy installed:\n'
29           r'   N:\georisk\downloads\numpy')
30    tkinter_error.tkinter_error(msg)
31    print msg
32    sys.exit(1)
33
34try:
35    import wx
36except ImportError:
37    import tkinter_error
38    msg = ('Sorry, you have to install WxPython.\n'
39           'Get it from either:\n'
40           r'   N:\georisk\downloads\event_selection, or' '\n'
41           '   [http://www.wxpython.org/download.php#binaries].')
42    tkinter_error.tkinter_error(msg)
43    print msg
44    sys.exit(1)
45
46Imported_PyEmbeddedImage = True
47try:
48    from wx.lib.embeddedimage import PyEmbeddedImage
49except ImportError:
50    Imported_PyEmbeddedImage = False
51
52
53# program name and version
54APP_NAME = 'plotcsv'
55APP_VERSION = '0.2'
56
57# name of the configuration filename
58# GUI values are saved in this file for next time
59ConfigFilename = '%s.cfg' % APP_NAME
60
61# GUI definitions
62BUTTON_WIDTH = 100
63BUTTON_HEIGHT = 30
64
65MARGIN = 10
66
67LAB_CTRL_OFFSET = 5
68
69BOX_WIDTH = 400
70BOX_HEIGHT = 360
71BOX_CSV_WIDTH = 400
72BOX_CSV_HEIGHT = 265
73BOX_PLOT_WIDTH = 400
74BOX_PLOT_HEIGHT = 140
75
76TXT_CSVFILE_WIDTH = 390
77TXT_CSVFILE_HEIGHT = 200
78
79COLLAB_X_OFFSET = 25
80
81CHBOX_HEIGHT = 30
82CHBOX_WIDTH = 100
83
84FORM_WIDTH = BOX_WIDTH + 15
85FORM_HEIGHT = 450
86
87START_YOFFSET = 7
88OUT_TEXT_XOFFSET = 8
89OUT_TEXT_YOFFSET = 11
90OUTPUT_BOX_HEIGHT = 45
91
92GEN_YOFFSET = 15
93GEN_DELTAY = 15
94GEN_LABELXOFFSET = 5
95GEN_TEXTXOFFSET = BOX_WIDTH/2 + GEN_LABELXOFFSET + 10
96
97TEXTLIST_WIDTH = BOX_WIDTH - 10
98TEXTLIST_HEIGHT = 200
99
100CTL_WIDTH = BOX_WIDTH / 2
101CTL_HEIGHT = 22
102
103DOUBLE_BUTTON_OFFSET = 8
104
105COMBO_FUDGE = 2
106
107# Copies of various parameters
108Region = ''
109GaugeNumber = ''
110MinimumHeight = ''
111MaximumHeight = ''
112
113# Flag strings - keys in the 'options' dictionary
114X_DATACOL = 'x_datacol'
115Y_DATACOL = 'y_datacol'
116X_RANGE = 'x_range'
117Y_RANGE = 'y_range'
118FILENAME = 'filename'
119SIZE = 'size'
120TITLE = 'title'
121X_LABEL = 'x_label'
122Y_LABEL = 'y_label'
123
124CSVWildcard = 'CSV files (*.csv)|*.csv|All files (*.*)|*.*'
125
126# Flag strings - keys in the 'options' dictionary
127X_DATACOL = 'x_datacol'
128Y_DATACOL = 'y_datacol'
129X_RANGE = 'x_range'
130Y_RANGE = 'y_range'
131FILENAME = 'filename'
132SIZE = 'size'
133TITLE = 'title'
134X_LABEL = 'x_label'
135Y_LABEL = 'y_label'
136
137
138################################################################################
139# This code was generated by img2py.py
140# Embed the ICON image here, so we don't need an external *.ico file.
141################################################################################
142
143if Imported_PyEmbeddedImage:
144    def getIcon():
145        return PyEmbeddedImage(
146    "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAABHNCSVQICAgIfAhkiAAAApxJ"
147    "REFUWIXtlr9O40AQh2edWwsKJFAarGtcRaElJWkoo4TGZegocqLME6TOC6CtaWjCFXRUfgJE"
148    "xStsiivuKiQ4ND8Kx/bu+k+cUKRhpC0ie+f7sjs7a0FEoB3GDyIiYDcOQgjydkI24lvgW2Bj"
149    "gctLolaLaG+P6OiI6PCQSIh89HpEy+VmOdE0lAKI1o8oArSuycMKAQegVQ9qDJeymUCdhGIF"
150    "yRLE1FzAhZ+dATc3wN0d8PiYg9z3XAkLbgl0u4DnAePxWvja5XXen0zK4RFHhoA5lNoabs4z"
151    "Uw4finDN2hAIw/xtKQGltoanMZut5k4U6K0IB2AIaJ0QVrQPT+LaU1vD0xg+2PDuSwJXCggC"
152    "UwAoSLyTxC9SpfAFLxBwAJ99THlavg1uwS0i0LHGcGiurHMKbucav0Uu8V9I/JsrK7FmjRM+"
153    "yRL77CPmuBbefUngxeNqCKR7HpDGPUWFmkjhEUf5v1qNkMNMoqzaNWtMJsU+kQmMx/bDq4HG"
154    "68CW0LdzC97iFtrctiSmPC2Fpzvc7yenvd9PfmcC+/s56/x8tedGTehjQnQvssSSJRQrxBwj"
155    "5LCwIi68KqwV8LykHz0/mxuuoa8GiBZ5YvnhQXFeF2USTeCWQFVo1oheBzn8jaCuPatZpRId"
156    "7sBnHyMeJXCtgaen2vNbK+AWnHwXUJNiYZZP1slylt1Sp6eZVKVAAc4S6u/c6hO1Em7JV1yV"
157    "pQKl8HTPnWZVKpH14PX3dUGgFm4ub5WEewvNZvZc54KxBBrB6ySmU+DgoHgPu2FIZAIbwask"
158    "zNHrrf9oCIxPMrO3N4LXSYQhEMeNpmcCKVywaA43JUYjwPeBTqcx3BIIOYTHHi74YjP4F4OI"
159    "IIgI9JOI2kS0JKI/m33TfzU+ASG0bvT61fpzAAAAAElFTkSuQmCC")
160
161
162def error(msg):
163    print msg
164    import tkinter_error
165    tkinter_error.tkinter_error(msg)
166    sys.exit(10)
167
168
169################################################################################
170# An object that behaves like a dictionary but is pickled/unpickled from a file.
171# Used to save state in the application.
172################################################################################
173
174class Config(object):
175    """An object that behaves like a dictionary but is persistent:
176
177        cfg = Config('filename')
178        cfg['save'] = 'A saved value'
179        cfg[123] = 'value 123'
180        print cfg['save']
181
182    The values stored in the file 'filename' will be available to any
183    application using that config file.
184    """
185
186    def __init__(self, configfile=None):
187        """__init__(self, String filename=None) -> Config"""
188
189        self.delconf = False
190        self.cfgdict = {}
191        self.changed = False
192        if not configfile:
193            self.delconf = True
194            configfile =tempfile.mktemp()
195        self.configfile = os.path.abspath(configfile)
196        if os.path.exists(self.configfile):
197            try:
198                f = open(configfile, "r")
199                u = pickle.Unpickler(f)
200                self.cfgdict = u.load()
201            except pickle.UnpicklingError, e:
202                print e
203            except:
204                pass
205            else:
206                f.close()
207
208    def __setitem__(self, key, value):
209        """Override to allow: cfg[<key>] = <value>"""
210       
211        self.cfgdict[key] = value
212        self.changed = True
213
214    def __getitem__(self, key):
215        """Override to allow: <var> = cfg[<key>]"""
216       
217        return self.cfgdict.get(key, None)
218
219    def __str__(self):
220        """__str__(self) -> String"""
221       
222        return "<config object at %s>" % hex(id(self))
223
224    def getfilename(self):
225        """getfilename(self) -> String filename"""
226       
227        return self.configfile
228
229    def setdeleted(self):
230        """setdeleted(self)"""
231       
232        self.delconf = True
233
234    def save(self):
235        """save(self)"""
236       
237        try:
238            f = open(self.configfile, "w")
239            p = pickle.Pickler(f)
240            p.dump(self.cfgdict)
241        except pickle.PicklingError, e:
242            print e
243        else:
244            f.close()
245        self.changed = False
246
247    def close(self):
248        """close(self)"""
249       
250        if self.changed:
251            self.save()
252        if self.delconf:
253            if os.path.exists(self.configfile):
254                os.remove(self.configfile)
255
256    def __del__(self):
257        """__del__(self)"""
258       
259        self.close()
260       
261
262################################################################################
263# Plot routines
264################################################################################
265
266##
267# @brief Get CSV data from a file.
268# @param filename Path to the data file to plot.
269# @param x_hdr The X axis title string.
270# @param y_hdr The Y axis title string.
271def getCSVData(filename, x_hdr, y_hdr):
272    # get contents of data file
273    # after this, 'header' is list of column header strings
274    #             'data' is a list of lists of data
275    fd = open(filename)
276    c = csv.reader(fd)
277    data = []
278    for row in c:
279        data.append(row)
280    fd.close()
281    header = data[0]
282    del data[0]         # get rid of header in dataset
283
284    # ensure header strings don't have leading/trailing spaces
285    new_header = []
286    for h in header:
287        new_header.append(h.strip())
288    header = new_header
289
290    # get int index values for column headers
291    try:
292        x_index = header.index(x_hdr)
293    except ValueError:
294        error("Sorry, X column header '%s' isn't in data file '%s'."
295              % (x_hdr, filename))
296
297    try:
298        y_index = header.index(y_hdr)
299    except ValueError:
300        error("Sorry, Y column header '%s' isn't in data file '%s'."
301              % (y_hdr, filename))
302
303    # get appropriate columns from data[]
304    x_data = map(lambda x: float(x[x_index]), data)
305    if x_hdr == 'time':
306        x_data = map(lambda x: float(x)/3600., x_data)
307    y_data = map(lambda x: float(x[y_index]), data)
308
309    return (x_data, y_data)
310
311
312##
313# @brief Plot data files.
314# @param filenames List of full pathnames to plot.
315# @param x_hdr The X axis label string.
316# @param y_hdr The Y axis label string.
317# @param title The string used to title the graph.
318# @param legend True if a legend is to be displayed.
319# @param legend_path True if legend is to contain full file paths.
320def plot_files(filenames, x_hdr, y_hdr, title='', legend=False,
321               legend_path=False):
322    pylab.rc('axes', linewidth=2)
323    pylab.xlabel(x_hdr.title())
324    pylab.ylabel(y_hdr.title())
325    pylab.grid(True)
326
327    for f in filenames:
328        (x_data, y_data) = getCSVData(f, x_hdr, y_hdr)
329        pylab.plot(x_data, y_data)
330
331    if title:
332        pylab.title(title)
333    if legend:
334        (min_y, max_y) = pylab.ylim()
335        range_y = max_y - min_y
336        add_y = float(range_y) / 5
337        pylab.ylim((min_y, max_y+add_y))
338        legend_names = filenames
339        if not legend_path:
340            legend_names = [os.path.basename(x) for x in filenames]
341        pylab.legend(legend_names, 'upper right')
342
343    pylab.show()
344    pylab.close()
345
346
347################################################################################
348# GUI routines
349################################################################################
350
351class MyFrame(wx.Frame):
352    def __init__(self, parent, id, title, pos=wx.DefaultPosition,
353                    size=wx.DefaultSize, style=wx.DEFAULT_FRAME_STYLE):
354        '''Lay out the GUI form'''
355       
356        # Make the frame
357        wx.Frame.__init__(self, parent, id, title, pos=(50, 50),
358                            size=(FORM_WIDTH, FORM_HEIGHT),
359                            style=(wx.DEFAULT_FRAME_STYLE &
360                                   ~ (wx.RESIZE_BOX | wx.MAXIMIZE_BOX |
361                                      wx.RESIZE_BORDER)))
362
363        p = self.panel = wx.Panel(self, -1)
364
365        if Imported_PyEmbeddedImage:
366            tsunami = getIcon()
367            icon = tsunami.GetIcon()
368            self.SetIcon(icon)
369       
370        self.Center()
371        self.Show(True)
372
373        # start laying out controls
374        Y_OFFSET = START_YOFFSET
375        wx.StaticBox(p, -1, 'CSV files', (3, Y_OFFSET),
376                     size=(BOX_CSV_WIDTH, BOX_CSV_HEIGHT))
377        Y_OFFSET += GEN_DELTAY
378        self.txtCSVFiles = wx.ListBox(p, -1, pos=(8, Y_OFFSET),
379                                      size=(TXT_CSVFILE_WIDTH,
380                                            TXT_CSVFILE_HEIGHT))
381        self.CSVFiles = []
382        Y_OFFSET += TXT_CSVFILE_HEIGHT + MARGIN
383        x = FORM_WIDTH/2 - BUTTON_WIDTH - DOUBLE_BUTTON_OFFSET
384        self.btnAddCSVFile = wx.Button(p, label="Add", pos=(x, Y_OFFSET),
385                                     size=(100, BUTTON_HEIGHT))
386        x = FORM_WIDTH/2 + DOUBLE_BUTTON_OFFSET
387        self.btnDelCSVFile = wx.Button(p, label="Delete", pos=(x, Y_OFFSET),
388                                     size=(100, BUTTON_HEIGHT))
389        Y_OFFSET += BUTTON_HEIGHT + MARGIN
390
391        wx.StaticBox(p, -1, 'Plot files', (3, Y_OFFSET),
392                     size=(BOX_PLOT_WIDTH, BOX_PLOT_HEIGHT))
393        Y_OFFSET += GEN_DELTAY
394        wx.StaticText(p, -1, 'X-Column',
395                      pos=(COLLAB_X_OFFSET, Y_OFFSET+LAB_CTRL_OFFSET), style=wx.ALIGN_LEFT)
396        self.cbXColHdr = wx.Choice(p, -1, pos=(COLLAB_X_OFFSET+65, Y_OFFSET),
397                                   size=(80, -1))
398        self.XColHdr = []
399        wx.StaticText(p, -1, 'Y-Column',
400                      pos=(FORM_WIDTH/2+COLLAB_X_OFFSET,
401                           Y_OFFSET+LAB_CTRL_OFFSET), style=wx.ALIGN_LEFT)
402        x = FORM_WIDTH/2 + COLLAB_X_OFFSET
403        self.cbYColHdr = wx.Choice(p, -1, pos=(x+65, Y_OFFSET), size=(80, -1))
404        self.YColHdr = []
405        Y_OFFSET += GEN_DELTAY*2
406
407        self.chkLegend = wx.CheckBox(p, -1, " Show graph legend",
408                                     pos=(COLLAB_X_OFFSET+20, Y_OFFSET))       
409        self.chkLegendPath = wx.CheckBox(p, -1, " Full file path in legend",
410                                         pos=(COLLAB_X_OFFSET+180, Y_OFFSET))       
411        Y_OFFSET += GEN_DELTAY + 3
412       
413        wx.StaticText(p, -1, 'Title',
414                      pos=(18, Y_OFFSET+LAB_CTRL_OFFSET+2), style=wx.ALIGN_LEFT)
415        self.txtTitle = wx.TextCtrl(p, -1, "", size=(350, -1),
416                                    pos=(COLLAB_X_OFFSET+20,
417                                         Y_OFFSET+LAB_CTRL_OFFSET),
418                                    style=wx.ALIGN_LEFT)
419       
420        Y_OFFSET += BUTTON_HEIGHT + MARGIN
421        x = FORM_WIDTH/2 - BUTTON_WIDTH/2
422        self.btnPlot = wx.Button(p, label="Plot", pos=(x, Y_OFFSET),
423                                     size=(100, BUTTON_HEIGHT))
424
425        # bind controls/events to handlers
426        self.Bind(wx.EVT_BUTTON, self.AddCSVFile, self.btnAddCSVFile)
427        self.Bind(wx.EVT_BUTTON, self.DelCSVFile, self.btnDelCSVFile)
428        self.Bind(wx.EVT_BUTTON, self.PlotFiles, self.btnPlot)
429        self.Bind(wx.EVT_CHECKBOX, self.ChangeLegend, self.chkLegend)
430       
431        self.Bind(wx.EVT_CLOSE, self.doStateSave)
432
433        # restore saved values
434        self.common_headers = None
435        self.cfg = Config(ConfigFilename)
436        self.doStateRestore()
437        self.updateHdrChoices()
438       
439    def AddCSVFile(self, event):
440        """Add a CSV file to the listbox"""
441
442        dlg = wx.FileDialog(self, message='Choose a CSV file',
443                            defaultDir=self.file_dir, 
444                            defaultFile='',
445                            wildcard=CSVWildcard,
446                            style=wx.OPEN | wx.CHANGE_DIR)
447
448        if dlg.ShowModal() == wx.ID_OK:
449            # This returns a list of  files that were selected.
450            path = dlg.GetPath()
451            self.file_dir = dlg.GetDirectory()
452            dlg.Destroy()
453            if path not in self.CSVFiles:
454                headers = self.getHeaders(path)
455                if headers:
456                    common_headers = self.orHeaders(headers)
457                    if not common_headers or len(common_headers) < 2:
458                        self.error("Sorry, file '%s' doesn't have enough headers in common with current files" % path)
459                    else:
460                        self.headers.append(headers)
461                        self.common_headers = common_headers
462                        self.CSVFiles.append(path)
463                        self.txtCSVFiles.Append(path)
464                else:
465                    self.error("Sorry, file '%s' doesn't appear to be a CSV file" % path)
466                self.updateHdrChoices()
467
468    def DelCSVFile(self, event):
469        """Delete a CSV file from the listbox"""
470
471        sel = self.txtCSVFiles.GetSelections()
472        if sel:
473            sel = sel[0]
474            self.txtCSVFiles.Delete(sel)
475            del self.CSVFiles[sel]
476            del self.headers[sel]
477            self.common_headers = self.GenCommonHeaders(self.headers)
478            self.updateHdrChoices()
479
480    def ChangeLegend(self, event):
481        if self.chkLegend.GetValue():
482            self.chkLegendPath.Enable()
483        else:
484            self.chkLegendPath.Disable()
485
486    def doStateSave(self, event):
487        """Save state from 'self' variables & controls."""
488
489        self.cfg['CSVFiles'] = self.CSVFiles
490        self.cfg['XColHdr'] = self.XColHdr
491        self.cfg['XColHdrSelection'] = self.cbXColHdr.GetSelection()
492        self.cfg['YColHdr'] = self.YColHdr
493        self.cfg['YColHdrSelection'] = self.cbYColHdr.GetSelection()
494        self.cfg['file_dir'] = self.file_dir
495        self.cfg['GraphTitle'] = self.txtTitle.GetValue()
496        self.cfg['GraphLegend'] = self.chkLegend.GetValue()
497        self.cfg['LegendPath'] = self.chkLegendPath.GetValue()
498
499        self.cfg.save()
500        event.Skip(True)
501
502    def doStateRestore(self):
503        """Restore state to 'self' variables."""
504
505        self.CSVFiles = self.cfg['CSVFiles']
506        if self.CSVFiles is None:
507            self.CSVFiles = []
508        self.XColHdr = self.cfg['XColHdr']
509        if self.XColHdr is None:
510            self.XColHdr = []
511        self.XColHdrSelection = self.cfg['XColHdrSelection']
512        self.YColHdr = self.cfg['YColHdr']
513        if self.YColHdr is None:
514            self.YColHdr = []
515        self.YColHdrSelection = self.cfg['YColHdrSelection']
516        self.file_dir = self.cfg['file_dir']
517        if self.file_dir is None:
518            self.file_dir = os.getcwd()
519
520        # put data into controls
521        self.headers = []
522        for (i, f) in enumerate(self.CSVFiles):
523            headers = self.getHeaders(f)
524            self.headers.append(headers)
525            self.txtCSVFiles.Insert(f, i)
526        self.common_headers = self.GenCommonHeaders(self.headers)
527        self.updateHdrChoices()
528
529        if self.XColHdrSelection >= 0:
530            self.cbXColHdr.SetSelection(self.XColHdrSelection)
531        if self.YColHdrSelection >= 0:
532            self.cbYColHdr.SetSelection(self.YColHdrSelection)
533        title = self.cfg['GraphTitle']
534        if title is None:
535            title = ''
536        self.txtTitle.SetValue(title)
537        if self.cfg['GraphLegend']:
538            self.chkLegend.SetValue(True)
539        if self.cfg['LegendPath']:
540            self.chkLegendPath.SetValue(True)
541        self.ChangeLegend(None)
542       
543    def getHeaders(self, filename):
544        '''Get header list from a file'''
545
546        try:
547            fd = open(filename, 'r')
548            hdr = fd.readline()
549            fd.close()
550        except:
551            return None
552
553        hdr = hdr.split(',')
554        hdr = [x.strip() for x in hdr]
555
556        return hdr
557       
558    def updateHdrChoices(self):
559        '''Update choice controls with header lists.
560
561        Disable controls if no CSV files (ie, no headers).
562        Keep selections visible afterwards, if possible.
563        '''
564
565        # get current selections, if any
566        selected_x = self.cbXColHdr.GetStringSelection()
567        selected_y = self.cbYColHdr.GetStringSelection()
568       
569        self.cbXColHdr.Clear()
570        self.cbYColHdr.Clear()
571
572        if self.common_headers:
573            self.cbXColHdr.Enable()
574            self.cbYColHdr.Enable()
575            self.btnPlot.Enable()
576            for h in self.common_headers:
577                self.cbXColHdr.Append(h)
578                self.cbYColHdr.Append(h)
579            if selected_x in self.common_headers:
580                index = self.common_headers.index(selected_x)
581                self.cbXColHdr.SetSelection(index)
582            if selected_y in self.common_headers:
583                index = self.common_headers.index(selected_y)
584                self.cbYColHdr.SetSelection(index)
585        else:
586            self.cbXColHdr.Disable()
587            self.cbYColHdr.Disable()
588            self.btnPlot.Disable()
589
590    def orHeaders(self, new_header):
591        '''Update X & Y column header choices.
592
593        Return new 'common headers' list from current
594        self.common_headers and the supplied new_header list.
595        '''
596
597        if self.common_headers:
598            result = [x for x in new_header if x in self.common_headers]
599        else:
600            result = new_header
601        return result
602
603    def GenCommonHeaders(self, headers):
604        '''Get new set of common headers.
605
606        Return a new 'common header' list given a
607        list of header lists.
608        '''
609
610        result = []
611        for header in headers:
612            if result:
613                result = [x for x in result if x in header]
614            else:
615                result = header
616        return result
617
618    def PlotFiles(self, event):
619        '''Plot files in the CSV file listbox.'''
620       
621        selected_x = self.cbXColHdr.GetStringSelection()
622        selected_y = self.cbYColHdr.GetStringSelection()
623
624        if selected_x and selected_y:
625            plot_files(self.CSVFiles, selected_x, selected_y,
626                       title=self.txtTitle.GetValue(),
627                       legend=self.chkLegend.GetValue(),
628                       legend_path=self.chkLegendPath.GetValue()) 
629           
630    def error(self, msg):
631        '''Issue xwPython error message.'''
632       
633        dlg = wx.MessageDialog(self, msg, 'Error',
634                               wx.OK | wx.ICON_INFORMATION)
635        dlg.ShowModal()
636        dlg.Destroy()
637
638       
639################################################################################
640# Mainline code - start the application.
641################################################################################
642
643if __name__ == '__main__':
644    app = wx.App()
645    frame = MyFrame(None, -1, '%s %s' % (APP_NAME, APP_VERSION),
646                    size=(FORM_WIDTH, FORM_HEIGHT))
647    app.SetTopWindow(frame)
648    app.MainLoop()
649
Note: See TracBrowser for help on using the repository browser.