diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index d2619716..efcb1a8c 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -2,8 +2,6 @@ import os
 from datetime import datetime, timedelta
 from mimetypes import guess_type
 
-from telethon.tl.custom import Draft
-
 try:
     import socks
 except ImportError:
@@ -17,6 +15,7 @@ from .errors import (
 )
 from .network import ConnectionMode
 from .tl import TLObject
+from .tl.custom import Draft
 from .tl.entity_database import EntityDatabase
 from .tl.functions.account import (
     GetPasswordRequest
diff --git a/telethon/tl/custom/__init__.py b/telethon/tl/custom/__init__.py
index f960f3cc..40914f16 100644
--- a/telethon/tl/custom/__init__.py
+++ b/telethon/tl/custom/__init__.py
@@ -1 +1 @@
-from telethon.tl.custom.draft import Draft
+from .draft import Draft
diff --git a/telethon/tl/custom/draft.py b/telethon/tl/custom/draft.py
index 51affc21..c50baa78 100644
--- a/telethon/tl/custom/draft.py
+++ b/telethon/tl/custom/draft.py
@@ -1,15 +1,14 @@
-from telethon.tl.functions.messages import SaveDraftRequest
-from telethon.tl.types import UpdateDraftMessage
+from ..functions.messages import SaveDraftRequest
+from ..types import UpdateDraftMessage
 
 
 class Draft:
+    """
+    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 calling `client.get_drafts()`.
+    """
     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
 
@@ -22,8 +21,11 @@ class Draft:
     @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.")
+            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
@@ -36,14 +38,17 @@ class Draft:
 
     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::
+        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)
+            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
@@ -56,9 +61,10 @@ class Draft:
             message=text,
             no_webpage=no_webpage,
             reply_to_msg_id=reply_to_msg_id,
-            entities=entities))
+            entities=entities
+        ))
 
-        if result is True:
+        if result:
             self.text = text
             self.no_webpage = no_webpage
             self.reply_to_msg_id = reply_to_msg_id