Create RpcResult class and generalise core special cases

This results in a cleaner MTProtoSender, which now can always
read a TLObject with a guaranteed item, if the message is OK.

telethon/tl/core/__init__.py

@@ -0,0 +1,26 @@
+"""
+This module holds core "special" types, which are more convenient ways
+to do stuff in a `telethon.network.mtprotosender.MTProtoSender` instance.
+
+Only special cases are gzip-packed data, the response message (not a
+client message), the message container which references these messages
+and would otherwise conflict with the rest, and finally the RpcResult:
+
+    rpc_result#f35c6d01 req_msg_id:long result:bytes = RpcResult;
+
+Three things to note with this definition:
+1. The constructor ID is actually ``42d36c2c``.
+2. Those bytes are not read like the rest of bytes (length + payload).
+   They are actually the raw bytes of another object, which can't be
+   read directly because it depends on per-request information (since
+   some can return ``Vector<int>`` and ``Vector<long>``).
+3. Those bytes may be gzipped data, which needs to be treated early.
+"""
+from .tlmessage import TLMessage
+from .gzippacked import GzipPacked
+from .messagecontainer import MessageContainer
+from .rpcresult import RpcResult
+
+core_objects = {x.CONSTRUCTOR_ID: x for x in (
+    GzipPacked, MessageContainer, RpcResult
+)}

telethon/tl/core/gzippacked.py

@@ -0,0 +1,42 @@
+import gzip
+import struct
+
+from .. import TLObject
+
+
+class GzipPacked(TLObject):
+    CONSTRUCTOR_ID = 0x3072cfa1
+
+    def __init__(self, data):
+        super().__init__()
+        self.data = data
+
+    @staticmethod
+    def gzip_if_smaller(request):
+        """Calls bytes(request), and based on a certain threshold,
+           optionally gzips the resulting data. If the gzipped data is
+           smaller than the original byte array, this is returned instead.
+
+           Note that this only applies to content related requests.
+        """
+        data = bytes(request)
+        # TODO This threshold could be configurable
+        if request.content_related and len(data) > 512:
+            gzipped = bytes(GzipPacked(data))
+            return gzipped if len(gzipped) < len(data) else data
+        else:
+            return data
+
+    def __bytes__(self):
+        # TODO Maybe compress level could be an option
+        return struct.pack('<I', GzipPacked.CONSTRUCTOR_ID) + \
+               TLObject.serialize_bytes(gzip.compress(self.data))
+
+    @staticmethod
+    def read(reader):
+        assert reader.read_int(signed=False) == GzipPacked.CONSTRUCTOR_ID
+        return gzip.decompress(reader.tgread_bytes())
+
+    @classmethod
+    def from_reader(cls, reader):
+        return GzipPacked(gzip.decompress(reader.tgread_bytes()))

telethon/tl/core/messagecontainer.py

@@ -0,0 +1,53 @@
+import struct
+
+from ..tlobject import TLObject
+from .tlmessage import TLMessage
+
+
+class MessageContainer(TLObject):
+    CONSTRUCTOR_ID = 0x73f1f8dc
+
+    def __init__(self, messages):
+        super().__init__()
+        self.content_related = False
+        self.messages = messages
+
+    def to_dict(self, recursive=True):
+        return {
+            'content_related': self.content_related,
+            'messages':
+                ([] if self.messages is None else [
+                    None if x is None else x.to_dict() for x in self.messages
+                ]) if recursive else self.messages,
+        }
+
+    def __bytes__(self):
+        return struct.pack(
+            '<Ii', MessageContainer.CONSTRUCTOR_ID, len(self.messages)
+        ) + b''.join(bytes(m) for m in self.messages)
+
+    @staticmethod
+    def iter_read(reader):
+        reader.read_int(signed=False)  # code
+        size = reader.read_int()
+        for _ in range(size):
+            inner_msg_id = reader.read_long()
+            inner_sequence = reader.read_int()
+            inner_length = reader.read_int()
+            yield TLMessage(inner_msg_id, inner_sequence,
+                            body=reader.read(inner_length))
+
+    def __str__(self):
+        return TLObject.pretty_format(self)
+
+    def stringify(self):
+        return TLObject.pretty_format(self, indent=0)
+
+    @classmethod
+    def from_reader(cls, reader):
+        # This assumes that .read_* calls are done in the order they appear
+        return MessageContainer([TLMessage(
+            msg_id=reader.read_long(),
+            seq_no=reader.read_int(),
+            body=reader.read(reader.read_int())
+        ) for _ in range(reader.read_int())])

telethon/tl/core/rpcresult.py

@@ -0,0 +1,23 @@
+from .gzippacked import GzipPacked
+from ..types import RpcError
+
+
+class RpcResult:
+    CONSTRUCTOR_ID = 0xf35c6d01
+
+    def __init__(self, req_msg_id, body, error):
+        self.req_msg_id = req_msg_id
+        self.body = body
+        self.error = error
+
+    @classmethod
+    def from_reader(cls, reader):
+        msg_id = reader.read_long()
+        inner_code = reader.read_int(signed=False)
+        if inner_code == RpcError.CONSTRUCTOR_ID:
+            return RpcResult(msg_id, None, RpcError.from_reader(reader))
+        if inner_code == GzipPacked.CONSTRUCTOR_ID:
+            return RpcResult(msg_id, GzipPacked.from_reader(reader).data, None)
+
+        reader.seek(-4)
+        return RpcResult(msg_id, reader.read(), None)

telethon/tl/core/tlmessage.py

@@ -0,0 +1,68 @@
+import asyncio
+import struct
+
+from .gzippacked import GzipPacked
+from .. import TLObject
+from ..functions import InvokeAfterMsgRequest
+
+
+class TLMessage(TLObject):
+    """
+    https://core.telegram.org/mtproto/service_messages#simple-container.
+
+    Messages are what's ultimately sent to Telegram:
+        message msg_id:long seqno:int bytes:int body:bytes = Message;
+
+    Each message has its own unique identifier, and the body is simply
+    the serialized request that should be executed on the server. Then
+    Telegram will, at some point, respond with the result for this msg.
+
+    Thus it makes sense that requests and their result are bound to a
+    sent `TLMessage`, and this result can be represented as a `Future`
+    that will eventually be set with either a result, error or cancelled.
+    """
+    def __init__(self, msg_id, seq_no, body=None, request=None, after_id=0):
+        super().__init__()
+        self.msg_id = msg_id
+        self.seq_no = seq_no
+        self.container_msg_id = None
+        self.future = asyncio.Future()
+
+        # TODO Perhaps it's possible to merge body and request?
+        # We need things like rpc_result and gzip_packed to
+        # be readable by the ``BinaryReader`` for such purpose.
+
+        # Used for incoming, not-decoded messages
+        self.body = body
+
+        # Used for outgoing, not-encoded messages
+        self.request = request
+
+        # After which message ID this one should run. We do this so
+        # InvokeAfterMsgRequest is transparent to the user and we can
+        # easily invoke after while confirming the original request.
+        self.after_id = after_id
+
+    def to_dict(self, recursive=True):
+        return {
+            'msg_id': self.msg_id,
+            'seq_no': self.seq_no,
+            'request': self.request,
+            'container_msg_id': self.container_msg_id,
+            'after_id': self.after_id
+        }
+
+    def __bytes__(self):
+        if self.after_id is None:
+            body = GzipPacked.gzip_if_smaller(self.request)
+        else:
+            body = GzipPacked.gzip_if_smaller(
+                InvokeAfterMsgRequest(self.after_id, self.request))
+
+        return struct.pack('<qii', self.msg_id, self.seq_no, len(body)) + body
+
+    def __str__(self):
+        return TLObject.pretty_format(self)
+
+    def stringify(self):
+        return TLObject.pretty_format(self, indent=0)