From 0d2bd648715b851f2506b5716f98e0d331234344 Mon Sep 17 00:00:00 2001
From: KurimuzonAkuma <arturgayauskas@gmail.com>
Date: Fri, 24 May 2024 13:18:36 +0300
Subject: [PATCH] Handle FloodPremiumWait error

Closes KurimuzonAkuma/pyrogram#64
---
 compiler/errors/source/420_FLOOD.tsv | 1 +
 pyrogram/client.py                   | 7 ++++---
 pyrogram/session/session.py          | 4 ++--
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/compiler/errors/source/420_FLOOD.tsv b/compiler/errors/source/420_FLOOD.tsv
index b45fb3f6..f2969506 100644
--- a/compiler/errors/source/420_FLOOD.tsv
+++ b/compiler/errors/source/420_FLOOD.tsv
@@ -1,5 +1,6 @@
 id	message
 2FA_CONFIRM_WAIT_X	A wait of {value} seconds is required because this account is active and protected by a 2FA password
+FLOOD_PREMIUM_WAIT_X	A wait of {value} seconds is required
 FLOOD_TEST_PHONE_WAIT_X	A wait of {value} seconds is required in the test servers
 FLOOD_WAIT_X	A wait of {value} seconds is required
 PREMIUM_SUB_ACTIVE_UNTIL_X	A wait of {value} seconds is required
diff --git a/pyrogram/client.py b/pyrogram/client.py
index b9d790be..a62bbd70 100644
--- a/pyrogram/client.py
+++ b/pyrogram/client.py
@@ -44,7 +44,8 @@ from pyrogram.errors import CDNFileHashMismatch
 from pyrogram.errors import (
     SessionPasswordNeeded,
     VolumeLocNotFound, ChannelPrivate,
-    BadRequest, AuthBytesInvalid
+    BadRequest, AuthBytesInvalid,
+    FloodWait, FloodPremiumWait
 )
 from pyrogram.handlers.handler import Handler
 from pyrogram.methods import Methods
@@ -863,7 +864,7 @@ class Client(Methods):
             if isinstance(e, asyncio.CancelledError):
                 raise e
 
-            if isinstance(e, pyrogram.errors.FloodWait):
+            if isinstance(e, (FloodWait, FloodPremiumWait)):
                 raise e
 
             return None
@@ -1097,7 +1098,7 @@ class Client(Methods):
                         await cdn_session.stop()
             except pyrogram.StopTransmission:
                 raise
-            except pyrogram.errors.FloodWait:
+            except (FloodWait, FloodPremiumWait):
                 raise
             except Exception as e:
                 log.exception(e)
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 013ed295..4df65f53 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -28,7 +28,7 @@ from pyrogram import raw
 from pyrogram.connection import Connection
 from pyrogram.crypto import mtproto
 from pyrogram.errors import (
-    RPCError, InternalServerError, AuthKeyDuplicated, FloodWait, ServiceUnavailable, BadMsgNotification,
+    RPCError, InternalServerError, AuthKeyDuplicated, FloodWait, FloodPremiumWait, ServiceUnavailable, BadMsgNotification,
     SecurityCheckMismatch
 )
 from pyrogram.raw.all import layer
@@ -397,7 +397,7 @@ class Session:
         while True:
             try:
                 return await self.send(query, timeout=timeout)
-            except FloodWait as e:
+            except (FloodWait, FloodPremiumWait) as e:
                 amount = e.value
 
                 if amount > sleep_threshold >= 0: