From 6df9fc558ef8f1a9cf7ae8b4fe59b96b20832317 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Wed, 27 Sep 2017 13:46:53 +0200
Subject: [PATCH] Create and use a new GzipPacked class, also when sending

---
 telethon/network/mtproto_sender.py | 12 ++++-----
 telethon/tl/__init__.py            |  1 +
 telethon/tl/gzip_packed.py         | 40 ++++++++++++++++++++++++++++++
 telethon/tl/message_container.py   |  5 ++--
 4 files changed, 49 insertions(+), 9 deletions(-)
 create mode 100644 telethon/tl/gzip_packed.py

diff --git a/telethon/network/mtproto_sender.py b/telethon/network/mtproto_sender.py
index 43c870f5..a63e936e 100644
--- a/telethon/network/mtproto_sender.py
+++ b/telethon/network/mtproto_sender.py
@@ -9,7 +9,7 @@ from ..errors import (
     rpc_message_to_error
 )
 from ..extensions import BinaryReader, BinaryWriter
-from ..tl import MessageContainer
+from ..tl import MessageContainer, GzipPacked
 from ..tl.all_tlobjects import tlobjects
 from ..tl.types import MsgsAck
 
@@ -68,10 +68,12 @@ class MtProtoSender:
         self._pending_receive.extend(requests)
         if len(requests) == 1:
             request = requests[0]
+            data = GzipPacked.gzip_if_smaller(request)
         else:
             request = MessageContainer(self.session, requests)
+            data = request.to_bytes()
 
-        self._send_packet(request.to_bytes(), request)
+        self._send_packet(data, request)
 
     def _send_acknowledges(self):
         """Sends a messages acknowledge for all those who _need_confirmation"""
@@ -376,11 +378,7 @@ class MtProtoSender:
 
     def _handle_gzip_packed(self, msg_id, sequence, reader, state):
         self._logger.debug('Handling gzip packed data')
-        reader.read_int(signed=False)  # code
-        packed_data = reader.tgread_bytes()
-        unpacked_data = gzip.decompress(packed_data)
-
-        with BinaryReader(unpacked_data) as compressed_reader:
+        with BinaryReader(GzipPacked.read(reader)) as compressed_reader:
             return self._process_msg(msg_id, sequence, compressed_reader, state)
 
     # endregion
diff --git a/telethon/tl/__init__.py b/telethon/tl/__init__.py
index e022408b..abd27a35 100644
--- a/telethon/tl/__init__.py
+++ b/telethon/tl/__init__.py
@@ -1,3 +1,4 @@
 from .tlobject import TLObject
 from .session import Session
+from .gzip_packed import GzipPacked
 from .message_container import MessageContainer
diff --git a/telethon/tl/gzip_packed.py b/telethon/tl/gzip_packed.py
new file mode 100644
index 00000000..c2fe17c7
--- /dev/null
+++ b/telethon/tl/gzip_packed.py
@@ -0,0 +1,40 @@
+import gzip
+
+from . import TLObject
+from ..extensions import BinaryWriter
+
+
+class GzipPacked(TLObject):
+    constructor_id = 0x3072cfa1
+
+    def __init__(self, data):
+        super().__init__()
+        self.data = data
+
+    @staticmethod
+    def gzip_if_smaller(request):
+        """Calls request.to_bytes(), 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 = request.to_bytes()
+        # TODO This threshold could be configurable
+        if request.content_related and len(data) > 512:
+            gzipped = GzipPacked(data).to_bytes()
+            return gzipped if len(gzipped) < len(data) else data
+        else:
+            return data
+
+    def to_bytes(self):
+        # TODO Maybe compress level could be an option
+        with BinaryWriter() as writer:
+            writer.write_int(GzipPacked.constructor_id, signed=False)
+            writer.tgwrite_bytes(gzip.compress(self.data))
+            return writer.get_bytes()
+
+    @staticmethod
+    def read(reader):
+        reader.read_int(signed=False)  # code
+        return gzip.decompress(reader.tgread_bytes())
diff --git a/telethon/tl/message_container.py b/telethon/tl/message_container.py
index ce7be98d..7686773b 100644
--- a/telethon/tl/message_container.py
+++ b/telethon/tl/message_container.py
@@ -1,4 +1,4 @@
-from . import TLObject
+from . import TLObject, GzipPacked
 from ..extensions import BinaryWriter
 
 
@@ -24,7 +24,8 @@ class MessageContainer(TLObject):
             writer.write_int(
                 self.session.generate_sequence(x.content_related)
             )
-            packet = x.to_bytes()
+
+            packet = GzipPacked.gzip_if_smaller(x)
             writer.write_int(len(packet))
             writer.write(packet)