From 2eb0aeba0dd93996f326ded13f7459b232ff4363 Mon Sep 17 00:00:00 2001
From: Markus Zolliker <markus.zolliker@psi.ch>
Date: Tue, 28 Jul 2020 14:23:33 +0200
Subject: [PATCH] add readbytes method to AsynConn

+ flush_recv now also clears _rxbuffer

Change-Id: I33c7ea1a9a1d8b663e5cd3bd81cf7ad43448e0fa
Reviewed-on: https://forge.frm2.tum.de/review/c/sine2020/secop/playground/+/23548
Tested-by: Jenkins Automated Tests <pedersen+jenkins@frm2.tum.de>
Reviewed-by: Enrico Faulhaber <enrico.faulhaber@frm2.tum.de>
Reviewed-by: Markus Zolliker <markus.zolliker@psi.ch>
---
 secop/lib/asynconn.py | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/secop/lib/asynconn.py b/secop/lib/asynconn.py
index 2bcb1ff..97be5c4 100644
--- a/secop/lib/asynconn.py
+++ b/secop/lib/asynconn.py
@@ -125,6 +125,28 @@ class AsynConn:
                 return None
             self._rxbuffer += data
 
+    def readbytes(self, nbytes, timeout=None):
+        """read one line
+
+        return either <nbytes> bytes or None if not enough data available within 1 sec (self.timeout)
+        if a non-zero timeout is given, a timeout error is raised instead of returning None
+        the timeout effectively used will not be lower than self.timeout (1 sec)
+        """
+        if timeout:
+            end = time.time() + timeout
+        while len(self._rxbuffer) < nbytes:
+            data = self.recv()
+            if not data:
+                if timeout:
+                    if time.time() < end:
+                        continue
+                    raise TimeoutError('timeout in readbytes (%g sec)' % timeout)
+                return None
+            self._rxbuffer += data
+        line = self._rxbuffer[:nbytes]
+        self._rxbuffer = self._rxbuffer[nbytes:]
+        return line
+
     def writeline(self, line):
         self.send(line + self.end_of_line)
 
@@ -154,9 +176,10 @@ class AsynTcp(AsynConn):
 
     def flush_recv(self):
         """flush recv buffer"""
-        data = []
+        data = [self._rxbuffer]
         while select.select([self.connection], [], [], 0)[0]:
             data.append(self.recv())
+        self._rxbuffer = b''
         return b''.join(data)
 
     def recv(self):
@@ -244,7 +267,9 @@ class AsynSerial(AsynConn):
         self.connection.write(data)
 
     def flush_recv(self):
-        return self.connection.read(self.connection.in_waiting)
+        result = self._rxbuffer + self.connection.read(self.connection.in_waiting)
+        self._rxbuffer = b''
+        return result
 
     def recv(self):
         """return bytes received within 1 sec"""