From e80eebc23494ecd558a4067fd0b6e3176cbc0261 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Mon, 6 May 2019 16:44:50 +0200
Subject: [PATCH] Add get_history_count method

---
 docs/source/pyrogram/Client.rst               |  1 +
 pyrogram/client/methods/messages/__init__.py  |  6 +-
 .../methods/messages/get_history_count.py     | 84 +++++++++++++++++++
 3 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 pyrogram/client/methods/messages/get_history_count.py

diff --git a/docs/source/pyrogram/Client.rst b/docs/source/pyrogram/Client.rst
index 3a7c231f..59293273 100644
--- a/docs/source/pyrogram/Client.rst
+++ b/docs/source/pyrogram/Client.rst
@@ -66,6 +66,7 @@ Messages
     delete_messages
     get_messages
     get_history
+    get_history_count
     iter_history
     send_poll
     vote_poll
diff --git a/pyrogram/client/methods/messages/__init__.py b/pyrogram/client/methods/messages/__init__.py
index d26c51cc..e843aa7c 100644
--- a/pyrogram/client/methods/messages/__init__.py
+++ b/pyrogram/client/methods/messages/__init__.py
@@ -16,7 +16,6 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from .stop_poll import StopPoll
 from .delete_messages import DeleteMessages
 from .download_media import DownloadMedia
 from .edit_message_caption import EditMessageCaption
@@ -25,6 +24,7 @@ from .edit_message_reply_markup import EditMessageReplyMarkup
 from .edit_message_text import EditMessageText
 from .forward_messages import ForwardMessages
 from .get_history import GetHistory
+from .get_history_count import GetHistoryCount
 from .get_messages import GetMessages
 from .iter_history import IterHistory
 from .retract_vote import RetractVote
@@ -44,6 +44,7 @@ from .send_venue import SendVenue
 from .send_video import SendVideo
 from .send_video_note import SendVideoNote
 from .send_voice import SendVoice
+from .stop_poll import StopPoll
 from .vote_poll import VotePoll
 
 
@@ -76,6 +77,7 @@ class Messages(
     RetractVote,
     DownloadMedia,
     IterHistory,
-    SendCachedMedia
+    SendCachedMedia,
+    GetHistoryCount
 ):
     pass
diff --git a/pyrogram/client/methods/messages/get_history_count.py b/pyrogram/client/methods/messages/get_history_count.py
new file mode 100644
index 00000000..046ec095
--- /dev/null
+++ b/pyrogram/client/methods/messages/get_history_count.py
@@ -0,0 +1,84 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+import logging
+import time
+from typing import Union
+
+from pyrogram.api import types, functions
+from pyrogram.client.ext import BaseClient
+from pyrogram.errors import FloodWait
+
+log = logging.getLogger(__name__)
+
+
+class GetHistoryCount(BaseClient):
+    def get_history_count(
+        self,
+        chat_id: Union[int, str]
+    ) -> int:
+        """Use this method to get the total count of messages in a chat.
+
+        .. note::
+
+            Due to Telegram latest internal changes, the server can't reliably find anymore the total count of messages
+            a **private** or a **basic group** chat has with a single method call. To overcome this limitation, Pyrogram
+            has to iterate over all the messages. Channels and supergroups are not affected by this limitation.
+
+        Args:
+            chat_id (``int`` | ``str``):
+                Unique identifier (int) or username (str) of the target chat.
+
+        Returns:
+            On success, an integer is returned.
+
+        Raises:
+            :class:`RPCError <pyrogram.RPCError>` in case of a Telegram RPC error.
+        """
+
+        peer = self.resolve_peer(chat_id)
+
+        if not isinstance(peer, types.InputPeerChannel):
+            offset = 0
+            limit = 100
+
+            while True:
+                try:
+                    r = self.send(
+                        functions.messages.GetHistory(
+                            peer=peer,
+                            offset_id=1,
+                            offset_date=0,
+                            add_offset=-offset - limit,
+                            limit=limit,
+                            max_id=0,
+                            min_id=0,
+                            hash=0
+                        )
+                    )
+                except FloodWait as e:
+                    log.warning("Sleeping for {}s".format(e.x))
+                    time.sleep(e.x)
+                    continue
+
+                if not r.messages:
+                    return offset
+
+                offset += len(r.messages)
+
+        return self.get_history(chat_id=chat_id, limit=1).total_count