support writing files for plot3d

1 files changed, 9 insertions, 1 deletions

diff --git a/util.py b/util.py
index 480ae5f..7d81d4f 100644
--- a/util.py
+++ b/util.py
@@ -8,7 +8,8 @@ def load_csv_file(fname, *args, **kargs):
     f.close() # don't leak file handles
     return data
 
-def write_csv_file(fname, rows, header=None, width=None):
+def write_csv_file(fname, rows, header=None, width=None,
+                   break_col=None):
     if fname is None:
         f = Tmp()
     else:
@@ -22,10 +23,17 @@ def write_csv_file(fname, rows, header=None, width=None):
         f.write('#')
         f.write(", ".join([fmt % str(x) for x in header]))
         f.write('\n')
+    last_row = None
     for row in rows:
+        # Insert extra line breaks for Gnuplot plot3d plotting whenever
+        # the value of break_col changes.
+        if not last_row is None and not break_col is None:
+            if last_row[break_col] != row[break_col]:
+                f.write('\n')
         f.write(' ')
         f.write(", ".join([fmt % str(x) for x in row]))
         f.write('\n')
+        last_row = row
     if fname is None:
         # keep it open
         f.flush()