1 files changed, 175 insertions, 4 deletions

diff --git a/gnuplot.py b/gnuplot.py
index 260fbf5..aa9e26b 100755
--- a/gnuplot.py
+++ b/gnuplot.py
@@ -6,6 +6,13 @@ from optparse import make_option as o
 FORMATS = ['png', 'eps', 'pdf', 'show', 'ps']
 STYLES  = ['points', 'lines', 'linespoints']
 
+TERMINALS = {
+    'png' : 'png',
+    'eps' : 'postscript eps',
+    'ps'  : 'postscript',
+    'pdf' : 'pdf',
+    }
+
 class CommandBuffer(object):
     def __init__(self):
         self.cmds = []
@@ -16,6 +23,21 @@ class CommandBuffer(object):
     def __str__(self):
         return '\n'.join([str(x) for x in self.cmds])
 
+
+class Label(object):
+    def __init__(self, x, y, txt):
+        self.x   = x
+        self.y   = y
+        self.txt = txt
+        # support more label stuff optionally...
+
+    def gnuplot_cmd(self):
+        return '"%s" at %f,%f center' % (self.txt, self.x, self.y)
+
+    def __str__(self):
+        return 'set label %s' % self.gnuplot_cmd()
+
+
 class FileGraph(object):
     def __init__(self, fname, xcol=1, ycol=2, error=None,
                  title=None,
@@ -25,12 +47,12 @@ class FileGraph(object):
         self.ycol  = ycol
         self.error = error
         self.title = title
-        self.style = style
+        self.style = style # replace with linetype, lines/points etc.
 
     def __str__(self):
-        return self.get_command()
+        return self.gnuplot_cmd()
 
-    def get_command(self, default_style=None):
+    def gnuplot_cmd(self, default_style=None):
         if self.error:
             using_txt = "%s:%s:%s" % (self.xcol, self.ycol, self.error)
             style     = self.style if self.style else "errorbars"
@@ -42,6 +64,155 @@ class FileGraph(object):
         return "'%s' using %s%s%s" % \
             (self.fname, using_txt, title_txt, style_txt)
 
+
+class LiteralGraph(object):
+    def __init__(self, literal_expr, title=None, style=None):
+        self.expr  = literal_expr
+        self.title = title
+        self.style = style
+
+    def __str__(self):
+        return self.gnuplot_cmd()
+
+    def gnuplot_cmd(self, default_style=None):
+        style     = self.style if self.style else default_style
+        style_txt = " with %s" % style if style else ""
+        title_txt = " title '%s'" % self.title if self.title else ""
+        return "%s %s%s" % \
+            (self.expr, title_txt, style_txt)
+
+
+label = Label
+
+def curve(fname=None, literal=None, **kargs):
+    if fname:
+        return FileGraph(fname, **kargs)
+    elif literal:
+        return LiteralGraph(literal, **kargs)
+
+
+class Plot(object):
+    def __init__(self):
+        self.format  = None
+        self.output  = None
+        self.size    = None
+
+        self.font      = None
+        self.font_size = None
+        self.enhanced_text = None
+
+        self.monochrome   = None
+        self.dashed_lines = None
+        self.rounded_caps = None
+
+        self.xrange = None
+        self.xticks = None
+        self.xlog   = None
+        self.xlabel = None
+
+        self.yrange = None
+        self.yticks = None
+        self.ylog   = None
+        self.ylabel = None
+
+        self.key    = None
+        self.title  = None
+
+        self.labels = []
+        self.curves = []
+
+    def gnuplot_commands(self, cmd_buf=None):
+        if cmd_buf:
+            g = cmd_buf
+        else:
+            g = CommandBuffer()
+
+        def isset(x):
+            return not (x is None)
+
+        if isset(self.format):
+            term = [TERMINALS[self.format]]
+        else:
+            term = 'x11'
+
+        ps_like = self.format in ['eps', 'pdf', 'ps']
+
+        if isset(self.monochrome) and ps_like:
+            term += ['monochrome' if self.monochrome else 'color']
+
+        if isset(self.enhanced_text):
+            term += ['enhanced' if self.enhanced_text else 'noenhanced']
+
+        if isset(self.dashed_lines) and ps_like:
+            term += ['dashed' if self.dashed_lines else 'solid']
+        if isset(self.font) and ps_like:
+            term += ['font "%s%s"' %
+                     (self.font,
+                      (",%s" % self.font_size) if isset(self.font_size)
+                      else "")
+                     ]
+
+        if isset(self.rounded_caps):
+            term += ['rounded' if self.rounded_caps else 'butt']
+
+        if isset(self.size):
+            term += ['size %s,%s' % (self.size)]
+
+        if isset(self.font_size) and self.format == 'png':
+            term += [self.font_size]
+
+        g("set terminal %s" % " ".join(term))
+        if isset(self.output):
+            g("set out '/dev/null'")
+
+        if self.xrange:
+            g("set xrange [%s:%s]" % self.xrange)
+        if self.xticks:
+            g("set xtics %s, %s" % self.xticks)
+        if self.xlabel:
+            g("set xlabel '%s'" % self.xlabel)
+
+        if self.yrange:
+            g("set yrange [%s:%s]" % self.yrange)
+        if self.yticks:
+            g("set ytics %s, %s" % self.yticks)
+        if self.ylabel:
+            g("set ylabel '%s'" % self.ylabel)
+
+        if self.key:
+            g('set key %s' % self.key)
+        if self.title:
+            g("set title '%s'" % self.title)
+
+        logscale = ""
+        if self.xlog:
+            logscale += "x"
+        if self.ylog:
+            logscale += "y"
+        if logscale:
+            g("set logscale %s" % logscale)
+
+        plots = [c.gnuplot_cmd() for c in self.curves]
+        if plots:
+            g("plot " + ", ".join(plots))
+
+        for l in self.labels:
+            g("set label " + l.gnuplot_cmd())
+
+        if isset(self.output):
+            g("set out '%s'" % self.output)
+
+        g("replot")
+
+        if isset(self.output):
+            g("set out") # close file
+
+        return g
+
+
+    def gnuplot_exec(self):
+        pipe2gnuplot(self.gnuplot_commands())
+
 def gnuplot_cmd(graphs, title=None, ylabel=None, xlabel=None,
                 format='show', term_opts=None,
                 style='linespoints', xrange=None,
@@ -99,7 +270,7 @@ def gnuplot_cmd(graphs, title=None, ylabel=None, xlabel=None,
             plot.append(str(gr))
         elif type(gr) == FileGraph:
             # formatter object
-            plot.append(gr.get_command(style))
+            plot.append(gr.gnuplot_cmd(style))
         elif len(gr) == 4:
             par = (gr[0], gr[1], gr[2], gr[3], style)
             plot += ["'%s' using %s:%s title '%s' with %s" % par]