Reformat connection classes to accommodate proxy settings

delivrance · delivrance · commit e9f6bce5795e · 2018-01-16T16:33:28.000+01:00
diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py
@@ -56,7 +56,7 @@ def close(self):
 
     def send(self, data: bytes):
         with self.lock:
-            self.connection.send(data)
+            self.connection.sendall(data)
 
     def recv(self) -> bytes or None:
-        return self.connection.recv()
+        return self.connection.recvall()
diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py
@@ -28,12 +28,6 @@ class TCP(socks.socksocket):
     def __init__(self):
         super().__init__()
 
-    def send(self, *args):
-        pass
-
-    def recv(self, *args):
-        pass
-
     def close(self):
         try:
             self.shutdown(socket.SHUT_RDWR)
diff --git a/pyrogram/connection/transport/tcp/tcp_abridged.py b/pyrogram/connection/transport/tcp/tcp_abridged.py
@@ -33,7 +33,7 @@ def connect(self, address: tuple):
         self.is_first_packet = True
         log.info("Connected!")
 
-    def send(self, data: bytes):
+    def sendall(self, data: bytes, *args):
         length = len(data) // 4
 
         data = (
@@ -48,20 +48,20 @@ def send(self, data: bytes):
 
         super().sendall(data)
 
-    def recv(self) -> bytes or None:
-        length = self.recvall(1)
+    def recvall(self, length: int = 0) -> bytes or None:
+        length = super().recvall(1)
 
         if length is None:
             return None
 
         if length == b"\x7f":
-            length = self.recvall(3)
+            length = super().recvall(3)
 
             if length is None:
                 return None
 
         length = int.from_bytes(length, "little") * 4
 
-        packet = self.recvall(length)
+        packet = super().recvall(length)
 
         return packet
diff --git a/pyrogram/connection/transport/tcp/tcp_full.py b/pyrogram/connection/transport/tcp/tcp_full.py
@@ -35,21 +35,21 @@ def connect(self, address: tuple):
         self.seq_no = 0
         log.info("Connected!")
 
-    def send(self, data: bytes):
+    def sendall(self, data: bytes, *args):
         # 12 = packet_length (4), seq_no (4), crc32 (4) (at the end)
         data = pack("<II", len(data) + 12, self.seq_no) + data
         data += pack("<I", crc32(data))
         self.seq_no += 1
 
         super().sendall(data)
 
-    def recv(self) -> bytes or None:
-        length = self.recvall(4)
+    def recvall(self, length: int = 0) -> bytes or None:
+        length = super().recvall(4)
 
         if length is None:
             return None
 
-        packet = self.recvall(unpack("<I", length)[0] - 4)
+        packet = super().recvall(unpack("<I", length)[0] - 4)
 
         if packet is None:
             return None
diff --git a/pyrogram/connection/transport/tcp/tcp_intermediate.py b/pyrogram/connection/transport/tcp/tcp_intermediate.py
@@ -34,7 +34,7 @@ def connect(self, address: tuple):
         self.is_first_packet = True
         log.info("Connected!")
 
-    def send(self, data: bytes):
+    def sendall(self, data: bytes, *args):
         length = len(data)
         data = pack("<i", length) + data
 
@@ -44,12 +44,12 @@ def send(self, data: bytes):
 
         super().sendall(data)
 
-    def recv(self) -> bytes or None:
-        length = self.recvall(4)
+    def recvall(self, length: int = 0) -> bytes or None:
+        length = super().recvall(4)
 
         if length is None:
             return None
 
-        packet = self.recvall(unpack("<I", length)[0])
+        packet = super().recvall(unpack("<I", length)[0])
 
         return packet
