In a previous post I covered the basics of graphing in Python with the MatPlotLib module.  In this post I am going to demostrate how to use MatPlotLib with some real world data retrieved from a web service and then processed into a format usable by MatPLotLib.

The example script performs the following steps:

  1. Takes a specified stock’s ticker symbol and column to plot over time (from Open, High, Low, Close, Volume, Adj Close) as input.
  2. Fetches the corresponding stock data from Yahoo! Finance and saves it into a CSV file using the urllib module.
  3. Processes the data in the CSV file into a suitable format for matplotlib using the csv, datetime and matplotlib.dates modules.
  4. Plots a graph of the data plotted over time using MatPlotLib and a saves a copy as PNG format image.

Note: To keep the example concise I am not performing any error handling.

"""Fetches specified stock data from Yahoo and graph it with MatPlotLib."""

from urllib import urlretrieve
from csv import DictReader
from matplotlib import pyplot
from matplotlib.dates import date2num
from datetime import datetime

def fetchstockdata( stockticker, filename ):
    """Fetch specified stock data and store it in named file."""
    url = 'http://ichart.finance.yahoo.com/table.csv?s=%s' % stockticker
    urlretrieve( url, filename )

def importstockdata( filename ):
    """Import CSV data into dict of lists, converting dates into timestamps."""
    results = {}
    for row in DictReader( open( filename,'rb' ) ):
        for col in row.keys():
            if col == 'Date':
                coldata = date2num( datetime.strptime( row[col], '%Y-%m-%d') )
            else:
                coldata = row[col]
            results.setdefault( col, [] ).append( coldata )
    return results

def plotstockdata( stockdata, stockticker, dates, col ):
    """Use MatPlotLib to graph speciifed stock data."""
    pyplot.plot_date( stockdata[dates], stockdata[col], '-', xdate=True )
    pyplot.title( '%s - %s / %s' % (stockticker, col, dates) )
    pyplot.xlabel( dates )
    pyplot.ylabel( col )
    pyplot.savefig( '%s.png' % stockticker )
    pyplot.show()

if __name__ == '__main__':
    from sys import argv
    # Use second argument as ticker and third argument as column.
    TICKER = argv[1].upper()
    COL = argv[2]
    # Grab the stock data from Yahoo!
    FILENAME = '%s.csv' % TICKER
    fetchstockdata( TICKER, FILENAME )
    # Import the data.
    DATA = importstockdata( FILENAME )
    # Plot the graph with Date as X-Axis and User selected column as Y-Axis.
    plotstockdata( DATA, TICKER, 'Date', COL )

Running this script with using the command line “python StockChart.py goog ‘Adj Close'” will produce a chart like the following.

This is a good example of why I like Python’s batteries included philosophy so much: it means I spend more of my time writing interesting bits of code as the utility functionality I need has already been implemented or is only an easy_install away.