rrdweb/rrd.py
author Tero Marttila <terom@fixme.fi>
Tue, 25 Jan 2011 01:28:06 +0200
changeset 32 47e977c23ba2
parent 29 c756e522c9ac
permissions -rw-r--r--
implement rendering of pmacct rrd graphs, and a dir/top.png feature
"""
    Friendly wrapper around the rrdtool python interface
"""

import rrdtool

import tempfile
import logging
import os

log = logging.getLogger('rrdweb.rrd')

def normalize_option_key (key) :
    """
        Normalize the given option key.

        A -- is prepended, and _'s are converted to -
    """

    return '--' + str(key).replace('_', '-')

def normalize_option_multi (key, values) :
    """
        Normalize a list of option values, returning a series of --opt, val1, --opt, val2, ...
    """

    for value in values :
        yield key
        yield str(value)

def normalize_option (key, value) :
    """
        Normalize the given option to a series of cmd-args.

        If value is None or False, no cmd-args are emitted. If value is True, only --opt is emitted. If value is a list,
        --opt and value are emitted for each item.
        
        Otherwise, both --opt and str(value) are emitted.
    """

    key = normalize_option_key(key)

    if value is None or value is False :
        # omit
        return ()

    elif value is True :
        # flag
        return (key, )
    
    elif isinstance(value, list) :
        # list of option values
        return tuple(normalize_option_multi(key, value))

    else :
        # option value
        return (key, str(value))

def merge_opts (*all_opts) :
    """
        Merge the given series of opt dicts
    """

    out = dict()

    for opts in all_opts :
        # XXX: not strictly true, merge lists
        out.update(opts)

    return out

def run_cmd (func, pre_args, opts, post_args) :
    """
        Run the given rrdtool.* function, formatting the given positional arguments and options.
    """
    
    # series of (cmd-arg, cmd-arg, ...) tuples, giving all '--opt' and 'value' arguments for each keyword argument
    opt_items = (normalize_option(key, value) for key, value in opts.iteritems())

    # decomposed series of cmd-args for options
    opt_args = [item for items in opt_items for item in items]

    # positional arguments
    pre_args = [str(arg) for arg in pre_args]
    post_args = [str(arg) for arg in post_args]

    # full arguments
    args = pre_args + opt_args + ['--'] + post_args

    log.debug('rrdtool %s %s', func.__name__, args)

    return func(*args)

def graph (out, *args, **opts) :
    """
        Render a graph image and/or print a report from data stored in one or several RRDs.

        If the output path is given, the graph data is written to the given file, otherwise a temporary file is
        created. The resulting output file is returned as an opened file object, ready to read() the graph data.

        In the temporary file case, the returned file is unlinked before being returned.

        If out is explicitly given as False, no graph output is generated, and graph_file is returned as None.

        Returns a (width, height, [ report output lines ], graph_file) tuple.
    """
    
    # open/tempfile
    if out is None :
        # tempfile
        tmp_fd, out_path = tempfile.mkstemp('.png')

        # disregard the tmp_fd, we re-open the file after it's been written out
        os.close(tmp_fd)
        
    elif out is False :
        # no output
        tmp_fd = None
        out_path = ''

    else :
        # direct output
        tmp_fd = None
        out_path = out
    
    # render
    try :
        # invoke
        width, height, output = run_cmd(rrdtool.graph, (out_path, ), opts, args)
        
        # return resulting file
        if out_path :
            out_file = open(out_path, 'r')

        else :
            # no graph output
            out_file = None

        return width, height, output, out_file

    finally :
        if tmp_fd :
            # cleanup tempfile
            # XXX: definately not portable to windows
            os.unlink(out_path)

def create (rrd_path, *args, **opts) :
    """
        Set up a new Round Robin Database (RRD).
    """

    return run_cmd(rrdtool.create, (rrd_path, ), opts, args)

def update (rrd_path, *args, **opts) :
    """
        Store new data values into an RRD.
    """

    return run_cmd(rrdtool.update, (rrd_path, ), opts, args)