grapher.py

# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import sys
import os.path
from optparse import OptionParser

templates = {'timelog': "timelogtemplate.html",
             'snapshot': "snapshottemplate.html",
             'rss': "rsstemplate.html"}

def chooseGraph(graphtype, graphdata):
  """ Chooses a template and creates a unique file name for the graph.

  Required arguments:
  graphtype -- the type of graph to be created (timelog, snapshot, rss)
  graphdata -- a dictionary of values to be substituted in the template code

  Data format:
  graphtype - graphdata
  'rss' - {title, haxis, data: [[PID, private, shared], [#, value, value], ...}
  'snapshot' - {title, data: [['PID', memory type], [#, value], ...]}
  'timelog' - {title, vaxis, data: [['timestamp', PID#, ...], [...], ...}

  Returns: result from createGraph() or system exit

  """
  graphfile = graphtype + "-" + graphdata['timestamp'] + ".html"
  if graphtype in templates.keys():
    template = templates[graphtype]
  else:
    message = ("ERROR: " + graphtype +
               " is not a valid graph type (timelog, snapshot, rss)")
    sys.exit(message)
  return createGraph(template, graphfile, graphdata)

def createGraph(template, graphfile, graphdata):
  """ Creates a graph.

  Required arguments:
  template -- template file being used to create the graph
  graphfile -- file name of eventually completed HTML graph
  graphdata -- dictionary with values to be substituted in template

  Returns: graphfile or system exit

  """
  #TODO: add some check to make sure data is in the correct format?
  empty = getTemplate(template)
  filled = fillTemplate(empty, graphdata)
  success = saveGraph(filled, graphfile)
  if success:
    return "Graph can be found at " + graphfile
  else:
    sys.exit("ERROR: Failed to create graph.")

def getTemplate(template):
  """ Reads and returns the lines from file 'template'. """
  if os.path.exists(template):
    f = open(template, 'r')
    template_string = f.read()
    f.close()
    return template_string
  else:
    sys.exit("ERROR: Template is not in current directory or does not exist")

def fillTemplate(empty, data):
  """ Customizes graph template with provided data.

  Required arguments:
  empty -- template string containing variables needing substitution
  data -- dictionary providing variable to value mappings for substitution

  Returns:
  filled -- string with variables substituted with respective values

  """
  filled = empty % data
  return filled

def saveGraph(filled, graph):
  """ Writes to a file.

  Required arguments:
  graph -- file to be written to
  filled -- text to be written

  Returns:
  True -- if write is successful
  False -- if write is unsuccessful

  """
  try:
    f = open(graph, 'w')
    f.write(filled)
    f.close()
    return True
  except IOERROR:
    return False

def testgraph(argv=None):
  """
  Uses commandline and prepared substition dictionary to test functionality of
  this only this.
  """
  if argv is None:
    argv = sys.argv
  usage = "Usage: grapher.py graphtype data"
  parser = OptionParser(usage=usage)
  options, args = parser.parse_args()
  if len(args) < 2:
    print usage
  else:
    graphtype = args[0]
    variables = args[1]
    print createGraph(graphtype, variables)

if __name__=="__main__":
  testgraph()