recv function now accepts size, added a method to enable fastframe, commented code cleanup

1 files changed, 20 insertions, 45 deletions

diff --git a/tektronix.py b/tektronix.py
index 3af4e5c..3cd266d 100644
--- a/tektronix.py
+++ b/tektronix.py
@@ -56,26 +56,33 @@ class TekScope(object):
 
         self.sock.sendall(msg)
 
-    def recv(self):
+    def clear(self):
+        self.send('*cls')
+
+    def recv(self, size=None):
         buffer = ''
 
-        while True:
-            buffer += self.sock.recv(1024)
+        if size is None:
+            while True:
+                buffer += self.sock.recv(4096)
 
-            if buffer.endswith('\n'):
-                break
+                if buffer.endswith('\n'):
+                    break
+        else:
+            while len(buffer) < size:
+                buffer += self.sock.recv(min(size-len(buffer),4096))
 
         return buffer
 
     def query(self, msg):
         """Sends a query to the oscilloscope and returns the response."""
-        if not msg.endswith('\n'):
-            msg += '\n'
-
         self.send(msg)
-
         return self.recv().strip()
 
+    def enable_fastframe(self, count):
+        self.send('horizontal:fastframe:state 1')
+        self.send('horizontal:fastframe:count %i' % count)
+
     def set_sequence_mode(self):
         """Sets the oscilloscope in single acquisition mode."""
         self.send('acquire:stopafter sequence\n')
@@ -144,22 +151,6 @@ class TekScope(object):
 
         if dtype is None:
             dtype = get_dtype(self.get_preamble(channel))
-            # encoding = self.query('data:encdg?')[:3]
-
-            # if encoding == 'RIB':
-            #     dtype = '>i'
-            # elif encoding == 'RPB':
-            #     dtype = '>u'
-            # elif encoding == 'SRI':
-            #     dtype = '<i'
-            # elif encoding == 'SRP':
-            #     dtype = '<u'
-            # else:
-            #     raise ValueError('unkown encoding: %s' % encoding)
-
-            # dtype += self.query('data:width?').strip()
-
-
 
         self.send('curve?\n')
 
@@ -174,22 +165,13 @@ class TekScope(object):
 
         y = int(self.sock.recv(int(x[1])))
 
-        buffer = ''
-
-        while len(buffer) < y:
-            buffer += self.sock.recv(y-len(buffer))
-
-        waveform = np.fromstring(buffer,dtype)
+        waveform = np.fromstring(self.recv(y),dtype)
 
         # messages end with a newline
         eom = self.sock.recv(1024)
 
         if eom != '\n':
-            print 'x = ', x
-            print 'y = ', y
-            print eom
-
-        assert eom == '\n'
+            raise Exception("eom != '\n'")
 
         self.send(header)
 
@@ -220,13 +202,13 @@ if __name__ == '__main__':
 
     scope = TekScope(setup.scope_ip, setup.port)
 
+    scope.clear()
+
     root, ext = os.path.splitext(args[0])
 
     if ext == '':
         ext = '.hdf5'
 
-    #print (root, ext)
-
     for run in range(options.nruns):
         if options.nruns > 1:
             fileid = str(run).zfill(int(log10(options.nruns))+1)
@@ -235,8 +217,6 @@ if __name__ == '__main__':
 
         filename = root + fileid + ext
 
-        #print filename
-
         t0 = time.time()
 
         with h5py.File(filename, 'w') as f:
@@ -244,12 +224,9 @@ if __name__ == '__main__':
 
             active_channels = scope.get_active_channels()
 
-            dtypes = {}
             for channel in active_channels:
                 preamble = scope.get_preamble(channel)
 
-                #dtypes[channel] = get_dtype(preamble)
-
                 dataset = f.create_dataset('channel%i' % channel, (options.nevents, preamble['NR_PT']), dtype=get_dtype(preamble), chunks=(max(1,min(100, options.nevents//100)), preamble['NR_PT']), compression='gzip')
 
                 for key, value in preamble.iteritems():
@@ -263,8 +240,6 @@ if __name__ == '__main__':
             fastframe_state = int(scope.query('horizontal:fastframe:state?'))
             fastframe_count = int(scope.query('horizontal:fastframe:count?'))
 
-            #print 'fastframe_state = ', fastframe_state
-
             i = 0
             while i < options.nevents:
                 print '\rsaving event: %i' % (i+1),