added safety net

2 files changed, 13 insertions, 8 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 1c9e1d8..dccc2f7 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -1,4 +1,5 @@
 from qolab.file_utils import save_table_with_header
+
 class Trace:
     def __init__(self, label):
         self.label = label
@@ -22,9 +23,8 @@ class Trace:
         prefixed_header = [prefix+l for l in header]
         return prefixed_header
 
-    def write(self, fname):
-        pass
-        
+    def write(self, fname, item_format='e'):
+        save_table_with_header(fname, self.values, self.header(), item_format=item_format)
 
 
 class TraceXY:
@@ -97,11 +97,16 @@ class TraceSetSameX:
         
 
 if __name__ == '__main__': 
+    import numpy as np
+    print("Testing trace")
     x=Trace('x trace')
+    x.values = np.random.normal(2,2,(4,2))
     x.unit='s'
     x.tags['x_tag'] = 'xxxx'
     x.tags['x_tag2'] = 'xxxx'
+    x.write('xtrace.dat')
     y=Trace('y trace')
+    y.values = np.random.normal(2,2,(4,2))
     y.unit='V'
     y.tags['y_tag'] = 'yyyy'
     xy=TraceXY('xy trace')
diff --git a/qolab/file_utils/__init__.py b/qolab/file_utils/__init__.py
index af92c47..49c29ab 100644
--- a/qolab/file_utils/__init__.py
+++ b/qolab/file_utils/__init__.py
@@ -73,15 +73,15 @@ def save_table_with_header(fname, data, header='', comment_symbol='%', skip_head
     # itemFormat examples: 'e', '.15e', 'f'
     fname = filename2os_fname(fname)
     file_exist_flag = os.path.exists(fname)
-    r, c = data.shape
     item_format=str.join('', ['{', f':{item_format}', '}'])
     with open(fname, 'a') as f:
         if not skip_headers_if_file_exist:
             for l in header:
                 f.write(f'{comment_symbol} {l}\n')
-        for r in data:
-            l=item_separator.join( map(item_format.format, r))
-            f.write(l)
-            f.write('\n')
+        if data is not None:
+            for r in data:
+                l=item_separator.join( map(item_format.format, r))
+                f.write(l)
+                f.write('\n')
         f.close()