From 3844c943875fb89e229e1c82a63457f9997d5339 Mon Sep 17 00:00:00 2001
From: JosXa <joscha.goetzer@gmail.com>
Date: Mon, 9 Oct 2017 02:59:23 +0200
Subject: [PATCH] Added Draft object for convenience

Added new module for custom objects
Implemented get_drafts method on TelegramClient
---
 telethon/telegram_client.py    | 21 ++++++++--
 telethon/tl/custom/__init__.py |  1 +
 telethon/tl/custom/draft.py    | 74 ++++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 4 deletions(-)
 create mode 100644 telethon/tl/custom/__init__.py
 create mode 100644 telethon/tl/custom/draft.py

diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index ca359b47..d2619716 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -2,6 +2,8 @@ import os
 from datetime import datetime, timedelta
 from mimetypes import guess_type
 
+from telethon.tl.custom import Draft
+
 try:
     import socks
 except ImportError:
@@ -28,8 +30,8 @@ from .tl.functions.contacts import (
 )
 from .tl.functions.messages import (
     GetDialogsRequest, GetHistoryRequest, ReadHistoryRequest, SendMediaRequest,
-    SendMessageRequest, GetChatsRequest
-)
+    SendMessageRequest, GetChatsRequest,
+    GetAllDraftsRequest)
 
 from .tl.functions import channels
 from .tl.functions import messages
@@ -302,9 +304,20 @@ class TelegramClient(TelegramBareClient):
             [utils.find_user_or_chat(d.peer, entities, entities) for d in ds]
         )
 
-    # endregion
+    def get_drafts(self):  # TODO: Ability to provide a `filter`
+        """
+        Gets all open draft messages.
 
-    # region Message requests
+        Returns a list of custom `Draft` objects that are easy to work with: You can call
+        `draft.set_message('text')` to change the message, or delete it through `draft.delete()`.
+
+        :return List[telethon.tl.custom.Draft]: A list of open drafts
+        """
+        response = self(GetAllDraftsRequest())
+        self.session.process_entities(response)
+        self.session.generate_sequence(response.seq)
+        drafts = [Draft._from_update(self, u) for u in response.updates]
+        return drafts
 
     def send_message(self,
                      entity,
diff --git a/telethon/tl/custom/__init__.py b/telethon/tl/custom/__init__.py
new file mode 100644
index 00000000..f960f3cc
--- /dev/null
+++ b/telethon/tl/custom/__init__.py
@@ -0,0 +1 @@
+from telethon.tl.custom.draft import Draft
diff --git a/telethon/tl/custom/draft.py b/telethon/tl/custom/draft.py
new file mode 100644
index 00000000..1fb1aba5
--- /dev/null
+++ b/telethon/tl/custom/draft.py
@@ -0,0 +1,74 @@
+from telethon.tl.functions.messages import SaveDraftRequest
+from telethon.tl.types import UpdateDraftMessage
+
+
+class Draft:
+    def __init__(self, client, peer, draft):
+        """
+        A custom class that encapsulates a draft on the Telegram servers, providing an
+        abstraction to change the message conveniently. The library will return instances of this
+        class when executing `client.get_drafts`.
+        """
+
+        self._client = client
+        self._peer = peer
+
+        self.text = draft.message,
+        self.date = draft.date,
+        self.no_webpage = draft.no_webpage,
+        self.reply_to_msg_id = draft.reply_to_msg_id,
+        self.entities = draft.entities
+
+    @classmethod
+    def _from_update(cls, client, update):
+        if not isinstance(update, UpdateDraftMessage):
+            raise ValueError("You can only create a new `Draft` from a corresponding "
+                             "`UpdateDraftMessage` object.")
+        return cls(client=client, peer=update.peer, draft=update.draft)
+
+    @property
+    def entity(self):
+        return self._client.get_entity(self._peer)
+
+    @property
+    def input_entity(self):
+        return self._client.get_input_entity(self._peer)
+
+    def set_message(self, text, no_webpage=None, reply_to_msg_id=None, entities=None):
+        """
+        Changes the draft message on the Telegram servers. The changes are reflected in this
+        object. Changing only individual attributes like for example the `reply_to_msg_id` should be
+        done by providing the current values of this object, like so::
+
+            draft.set_message(draft.text,
+                              no_webpage=draft.no_webpage,
+                              reply_to_msg_id=NEW_VALUE,
+                              entities=draft.entities)
+
+        :param str text: New text of the draft
+        :param bool no_webpage: Whether to attach a web page preview
+        :param int reply_to_msg_id: Message id to reply to
+        :param list entities: A list of formatting entities
+        :return bool: `True` on success
+        """
+        result = self._client(SaveDraftRequest(
+            peer=self._peer,
+            message=text,
+            no_webpage=no_webpage,
+            reply_to_msg_id=reply_to_msg_id,
+            entities=entities))
+
+        if result is True:
+            self.text = text
+            self.no_webpage = no_webpage
+            self.reply_to_msg_id = reply_to_msg_id
+            self.entities = entities
+
+        return result
+
+    def delete(self):
+        """
+        Deletes this draft
+        :return bool: `True` on success
+        """
+        return self.set_message(text='')