From b22e3f9828d3443bb2a57feae350a4d9d120c481 Mon Sep 17 00:00:00 2001
From: Jason <thebosssg19@gmail.com>
Date: Thu, 7 Jul 2022 21:25:36 -0400
Subject: [PATCH] fixes streamlabs_polly ratelimit

---
 TTS/streamlabs_polly.py | 21 ++++++++------
 utils/voice.py          | 61 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/TTS/streamlabs_polly.py b/TTS/streamlabs_polly.py
index 1f8a039..b7365ab 100644
--- a/TTS/streamlabs_polly.py
+++ b/TTS/streamlabs_polly.py
@@ -2,6 +2,7 @@ import random
 import requests
 from requests.exceptions import JSONDecodeError
 from utils import settings
+from utils.voice import check_ratelimit
 
 voices = [
     "Brian",
@@ -42,16 +43,20 @@ class StreamlabsPolly:
             voice = str(settings.config["settings"]["tts"]["streamlabs_polly_voice"]).capitalize()
         body = {"voice": voice, "text": text, "service": "polly"}
         response = requests.post(self.url, data=body)
-        try:
-            voice_data = requests.get(response.json()["speak_url"])
-            with open(filepath, "wb") as f:
-                f.write(voice_data.content)
-        except (KeyError, JSONDecodeError):
+        if not check_ratelimit(response):
+            self.run(text, filepath, random_voice)
+
+        else:
             try:
-                if response.json()["error"] == "No text specified!":
-                    raise ValueError("Please specify a text to convert to speech.")
+                voice_data = requests.get(response.json()["speak_url"])
+                with open(filepath, "wb") as f:
+                    f.write(voice_data.content)
             except (KeyError, JSONDecodeError):
-                print("Error occurred calling Streamlabs Polly")
+                try:
+                    if response.json()["error"] == "No text specified!":
+                        raise ValueError("Please specify a text to convert to speech.")
+                except (KeyError, JSONDecodeError):
+                    print("Error occurred calling Streamlabs Polly")
 
     def randomvoice(self):
         return random.choice(self.voices)
diff --git a/utils/voice.py b/utils/voice.py
index c4f27bf..e92f0f4 100644
--- a/utils/voice.py
+++ b/utils/voice.py
@@ -1,4 +1,65 @@
 import re
+import sys
+from datetime import datetime
+import time as pytime
+from time import sleep
+
+from requests import Response
+
+if sys.version_info[0] >= 3:
+    from datetime import timezone
+
+
+def check_ratelimit(response: Response):
+    """
+    Checks if the response is a ratelimit response.
+    If it is, it sleeps for the time specified in the response.
+    """
+    if response.status_code == 429:
+        try:
+            time = int(response.headers["X-RateLimit-Reset"])
+            print(f"Ratelimit hit. Sleeping for {time - int(pytime.time())} seconds.")
+            sleep_until(time)
+            return False
+        except KeyError:  # if the header is not present, we don't know how long to wait
+            return False
+
+    return True
+
+
+def sleep_until(time):
+    """
+    Pause your program until a specific end time.
+    'time' is either a valid datetime object or unix timestamp in seconds (i.e. seconds since Unix epoch)
+    """
+    end = time
+
+    # Convert datetime to unix timestamp and adjust for locality
+    if isinstance(time, datetime):
+        # If we're on Python 3 and the user specified a timezone, convert to UTC and get tje timestamp.
+        if sys.version_info[0] >= 3 and time.tzinfo:
+            end = time.astimezone(timezone.utc).timestamp()
+        else:
+            zoneDiff = pytime.time() - (datetime.now() - datetime(1970, 1, 1)).total_seconds()
+            end = (time - datetime(1970, 1, 1)).total_seconds() + zoneDiff
+
+    # Type check
+    if not isinstance(end, (int, float)):
+        raise Exception('The time parameter is not a number or datetime object')
+
+    # Now we wait
+    while True:
+        now = pytime.time()
+        diff = end - now
+
+        #
+        # Time is up!
+        #
+        if diff <= 0:
+            break
+        else:
+            # 'logarithmic' sleeping to minimize loop iterations
+            sleep(diff / 2)
 
 
 def sanitize_text(text: str) -> str: