Make curses client capable of websocket _and_ raw tcp connections.

[plomrogue2-experiments] / new2 / plomrogue / io_tcp.py

diff --git a/new2/plomrogue/io_tcp.py b/new2/plomrogue/io_tcp.py
index 203021cd06b9b7c34e2eab8efa8749c1e6188155..b030f1b9f1c98332084763812bf152666004bb7a 100644 (file)
--- a/new2/plomrogue/io_tcp.py
+++ b/new2/plomrogue/io_tcp.py
@@ -6,6 +6,7 @@ socketserver.TCPServer.allow_reuse_address = True
 
 
 
+from plomrogue.errors import BrokenSocketConnection
 class PlomSocket:
 
     def __init__(self, socket):
@@ -32,7 +33,6 @@ class PlomSocket:
         <http://stackoverflow.com/q/34919846>
 
         """
-        from plomrogue.errors import BrokenSocketConnection
         escaped_message = ''
         for char in message:
             if char in ('\\', '$'):
@@ -46,6 +46,7 @@ class PlomSocket:
             try:
                 sent = self.socket.send(data[totalsent:])
                 socket_broken = sent == 0
+                totalsent = totalsent + sent
             except OSError as err:
                 if err.errno == 9:  # "Bad file descriptor", when connection broken
                     socket_broken = True
@@ -53,7 +54,6 @@ class PlomSocket:
                     raise err
             if socket_broken and not silent_connection_break:
                 raise BrokenSocketConnection
-            totalsent = totalsent + sent
 
     def recv(self):
         """Get full send()-prepared message from self.socket.
@@ -77,12 +77,14 @@ class PlomSocket:
         data = b''
         msg = b''
         while True:
-            data += self.socket.recv(1024)
+            try:
+                data = self.socket.recv(1024)
+            except OSError as err:
+                if err.errno == 9:  # "Bad file descriptor", when connection broken
+                    raise BrokenSocketConnection
             if 0 == len(data):
-                return
-            cut_off = 0
+                break
             for c in data:
-                cut_off += 1
                 if esc:
                     msg += bytes([c])
                     esc = False
@@ -93,7 +95,6 @@ class PlomSocket:
                         yield msg.decode()
                     except UnicodeDecodeError:
                         yield None
-                    data = data[cut_off:]
                     msg = b''
                 else:
                     msg += bytes([c])