Let all connection modes implement the new Connection

telethon/network/connection/tcpabridged.py

@@ -1,31 +1,43 @@
 import struct
 
-from .tcpfull import ConnectionTcpFull
+from .connection import Connection
 
 
-class ConnectionTcpAbridged(ConnectionTcpFull):
+class ConnectionTcpAbridged(Connection):
     """
     This is the mode with the lowest overhead, as it will
     only require 1 byte if the packet length is less than
     508 bytes (127 << 2, which is very common).
     """
-    async def connect(self, ip, port):
-        result = await super().connect(ip, port)
-        await self.conn.write(b'\xef')
-        return result
+    async def connect(self):
+        await super().connect()
+        await self.send(b'\xef')
 
-    async def recv(self):
-        length = struct.unpack('<B', await self.read(1))[0]
-        if length >= 127:
-            length = struct.unpack('<i', await self.read(3) + b'\0')[0]
+    def _write(self, data):
+        """
+        Define wrapper write methods for `TcpObfuscated` to override.
+        """
+        self._writer.write(data)
 
-        return await self.read(length << 2)
+    async def _read(self, n):
+        """
+        Define wrapper read methods for `TcpObfuscated` to override.
+        """
+        return await self._reader.readexactly(n)
 
-    async def send(self, message):
-        length = len(message) >> 2
+    def _send(self, data):
+        length = len(data) >> 2
         if length < 127:
             length = struct.pack('B', length)
         else:
             length = b'\x7f' + int.to_bytes(length, 3, 'little')
 
-        await self.write(length + message)
+        self._write(length + data)
+
+    async def _recv(self):
+        length = struct.unpack('<B', await self._read(1))[0]
+        if length >= 127:
+            length = struct.unpack(
+                '<i', await self._read(3) + b'\0')[0]
+
+        return await self._read(length << 2)