Improved TcpClient receive

main.py

@@ -9,7 +9,7 @@ if __name__ == '__main__':
 
     else:
         print('Loading interactive example...')
-        
+
         # First, initialize our TelegramClient and connect
         settings = load_settings()
         client = TelegramClient(session_user_id=settings.get('session_name', 'anonymous'),

network/tcp_client.py

@@ -1,5 +1,6 @@
 # Python rough implementation of a C# TCP client
 import socket
+from utils import BinaryWriter
 
 
 class TcpClient:
@@ -23,11 +24,12 @@ class TcpClient:
 
     def read(self, buffer_size):
         """Reads (receives) the specified bytes from the connected peer"""
-        # TODO improve (don't cast to list, use a mutable byte list instead or similar, see recv_into)
-        result = []
-        while len(result) < buffer_size:
-            left_data = buffer_size - len(result)
-            partial = self.socket.recv(left_data)
-            result.extend(list(partial))
+        with BinaryWriter() as writer:
+            while writer.written_count < buffer_size:
+                # When receiving from the socket, we may not receive all the data at once
+                # This is why we need to keep checking to make sure that we receive it all
+                left_count = buffer_size - writer.written_count
+                partial = self.socket.recv(left_count)
+                writer.write(partial)
 
-        return bytes(result)
+            return writer.get_bytes()

utils/binary_writer.py

@@ -14,6 +14,7 @@ class BinaryWriter:
 
         self.stream = stream
         self.writer = BufferedWriter(self.stream)
+        self.written_count = 0
 
     # region Writing
 
@@ -21,30 +22,37 @@ class BinaryWriter:
     def write_byte(self, value):
         """Writes a single byte value"""
         self.writer.write(pack('B', value))
+        self.written_count += 1
 
     def write_int(self, value, signed=True):
         """Writes an integer value (4 bytes), which can or cannot be signed"""
         self.writer.write(int.to_bytes(value, length=4, byteorder='little', signed=signed))
+        self.written_count += 4
 
     def write_long(self, value, signed=True):
         """Writes a long integer value (8 bytes), which can or cannot be signed"""
         self.writer.write(int.to_bytes(value, length=8, byteorder='little', signed=signed))
+        self.written_count += 8
 
     def write_float(self, value):
         """Writes a floating point value (4 bytes)"""
         self.writer.write(pack('<f', value))
+        self.written_count += 4
 
     def write_double(self, value):
         """Writes a floating point value (8 bytes)"""
         self.writer.write(pack('<d', value))
+        self.written_count += 8
 
     def write_large_int(self, value, bits, signed=True):
         """Writes a n-bits long integer value"""
         self.writer.write(int.to_bytes(value, length=bits // 8, byteorder='little', signed=signed))
+        self.written_count += bits // 8
 
     def write(self, data):
         """Writes the given bytes array"""
         self.writer.write(data)
+        self.written_count += len(data)
 
     # endregion
 
@@ -98,6 +106,11 @@ class BinaryWriter:
             self.writer.flush()
         return self.stream.getvalue()
 
+    def get_written_bytes_count(self):
+        """Gets the count of bytes written in the buffer.
+           This may NOT be equal to the stream length if one was provided when initializing the writer"""
+        return self.written_count
+
     # with block
     def __enter__(self):
         return self