From 0d09f9d8c1ee170ce555ea2d7e06e699ae38ce26 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Mon, 25 Jun 2018 11:34:10 +0200
Subject: [PATCH] Create a custom.Forward class

---
 readthedocs/telethon.tl.custom.rst |   8 ++
 telethon/tl/custom/__init__.py     |   1 +
 telethon/tl/custom/forward.py      | 123 +++++++++++++++++++++++++++++
 telethon/tl/custom/message.py      |  20 +++--
 4 files changed, 145 insertions(+), 7 deletions(-)
 create mode 100644 telethon/tl/custom/forward.py

diff --git a/readthedocs/telethon.tl.custom.rst b/readthedocs/telethon.tl.custom.rst
index 138811f6..e8b6c8e6 100644
--- a/readthedocs/telethon.tl.custom.rst
+++ b/readthedocs/telethon.tl.custom.rst
@@ -37,3 +37,11 @@ telethon\.tl\.custom\.messagebutton module
     :members:
     :undoc-members:
     :show-inheritance:
+
+telethon\.tl\.custom\.forward module
+------------------------------------
+
+.. automodule:: telethon.tl.custom.forward
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/telethon/tl/custom/__init__.py b/telethon/tl/custom/__init__.py
index 1d3f4068..8414d794 100644
--- a/telethon/tl/custom/__init__.py
+++ b/telethon/tl/custom/__init__.py
@@ -2,4 +2,5 @@ from .draft import Draft
 from .dialog import Dialog
 from .input_sized_file import InputSizedFile
 from .messagebutton import MessageButton
+from .forward import Forward
 from .message import Message
diff --git a/telethon/tl/custom/forward.py b/telethon/tl/custom/forward.py
new file mode 100644
index 00000000..737b2119
--- /dev/null
+++ b/telethon/tl/custom/forward.py
@@ -0,0 +1,123 @@
+from ...utils import get_input_peer
+
+
+class Forward:
+    """
+    Custom class that encapsulates a :tl:`MessageFwdHeader` providing an
+    abstraction to easily access information like the original sender.
+
+    Attributes:
+
+        original_fwd (:tl:`MessageFwdHeader`):
+            The original :tl:`MessageFwdHeader` instance.
+
+        Any other attribute:
+            Attributes not described here are the same as those available
+            in the original :tl:`MessageFwdHeader`.
+    """
+    def __init__(self, client, original, entities):
+        self.__dict__ = original.__dict__
+        self._client = client
+        self.original_fwd = original
+        self._sender = entities.get(original.from_id)
+        self._chat = entities.get(original.channel_id)
+
+        self._input_sender =\
+            get_input_peer(self._sender) if self._sender else None
+        self._input_chat =\
+            get_input_peer(self._chat) if self._chat else None
+
+    # TODO The pattern to get sender and chat is very similar
+    # and copy pasted in/to several places. Reuse the code.
+    #
+    # It could be an ABC with some ``resolve_sender`` abstract,
+    # so every subclass knew what tricks it can make to get
+    # the sender.
+
+    @property
+    def sender(self):
+        """
+        The :tl:`User` that sent the original message. This may be ``None``
+        if it couldn't be found or the message wasn't forwarded from an user
+        but instead was forwarded from e.g. a channel.
+        """
+        return self._sender
+
+    async def get_sender(self):
+        """
+        Returns `sender` but will make an API if necessary.
+        """
+        if not self.sender and self.original_fwd.from_id:
+            try:
+                self._sender = await self._client.get_entity(
+                    await self.get_input_sender())
+            except ValueError:
+                # TODO We could reload the message
+                pass
+
+        return self._sender
+
+    @property
+    def input_sender(self):
+        """
+        Returns the input version of `user`.
+        """
+        if not self._input_sender and self.original_fwd.from_id:
+            try:
+                self._input_sender = self._client.session.get_input_entity(
+                    self.original_fwd.from_id)
+            except ValueError:
+                pass
+
+        return self._input_sender
+
+    async def get_input_sender(self):
+        """
+        Returns `input_sender` but will make an API call if necessary.
+        """
+        # TODO We could reload the message
+        return self.input_sender
+
+    @property
+    def chat(self):
+        """
+        The :tl:`Channel` where the original message was sent. This may be
+        ``None`` if it couldn't be found or the message wasn't forwarded
+        from a channel but instead was forwarded from e.g. an user.
+        """
+        return self._chat
+
+    async def get_chat(self):
+        """
+        Returns `chat` but will make an API if necessary.
+        """
+        if not self.chat and self.original_fwd.channel_id:
+            try:
+                self._chat = await self._client.get_entity(
+                    await self.get_input_chat())
+            except ValueError:
+                # TODO We could reload the message
+                pass
+
+        return self._chat
+
+    @property
+    def input_chat(self):
+        """
+        Returns the input version of `chat`.
+        """
+        if not self._input_chat and self.original_fwd.channel_id:
+            try:
+                self._input_chat = self._client.session.get_input_entity(
+                    self.original_fwd.channel_id)
+            except ValueError:
+                pass
+
+        return self._input_chat
+
+    async def get_input_chat(self):
+        """
+        Returns `input_chat` but will make an API call if necessary.
+        """
+        # TODO We could reload the message
+        return self.input_chat
diff --git a/telethon/tl/custom/message.py b/telethon/tl/custom/message.py
index d728dc72..8e069300 100644
--- a/telethon/tl/custom/message.py
+++ b/telethon/tl/custom/message.py
@@ -1,6 +1,7 @@
 from .. import types
 from ...utils import get_input_peer, get_peer_id, get_inner_text
 from .messagebutton import MessageButton
+from .forward import Forward
 
 
 class Message:
@@ -52,14 +53,11 @@ class Message:
         if not self._input_chat and self._chat:
             self._input_chat = get_input_peer(self._chat)
 
-        self._fwd_from_entity = None
         if getattr(self.original_message, 'fwd_from', None):
-            fwd = self.original_message.fwd_from
-            if fwd.from_id:
-                self._fwd_from_entity = entities.get(fwd.from_id)
-            elif fwd.channel_id:
-                self._fwd_from_entity = entities.get(get_peer_id(
-                    types.PeerChannel(fwd.channel_id)))
+            self._forward = Forward(
+                self._client, self.original_message.fwd_from, entities)
+        else:
+            self._forward = None
 
     def __new__(cls, client, original, entities, input_chat):
         if isinstance(original, types.Message):
@@ -325,6 +323,14 @@ class Message:
         """True if the message is a reply to some other or not."""
         return bool(self.original_message.reply_to_msg_id)
 
+    @property
+    def forward(self):
+        """
+        Returns `telethon.tl.custom.forward.Forward` if the message
+        has been forwarded from somewhere else.
+        """
+        return self._forward
+
     def _set_buttons(self, sender, chat):
         """
         Helper methods to set the buttons given the input sender and chat.