From 036285b67372b535d1564d8260d3098158c91ced Mon Sep 17 00:00:00 2001
From: Riccardo Balbo <riccardo0blb@gmail.com>
Date: Tue, 19 Nov 2024 19:28:46 +0000
Subject: [PATCH] format and types

---
 nwcp.py | 37 ++++++++++++++++++-------------------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/nwcp.py b/nwcp.py
index 6f6e296..50f6ae7 100644
--- a/nwcp.py
+++ b/nwcp.py
@@ -16,6 +16,11 @@
 from loguru import logger
 
 
+class RateLimit:
+    backoff: int = 0
+    last_attempt_time: int = 0
+
+
 class MainSubscription:
     def __init__(self):
         self.requests_sub_id: Optional[str] = None
@@ -86,8 +91,8 @@ def __init__(self, private_key: Optional[str] = None, relay: Optional[str] = Non
 
         # Subscription
         self.sub = None
-        self.rate_limit = {}
-        
+        self.rate_limit: Dict[str, RateLimit] = {}
+
         # websocket connection
         self.ws = None
 
@@ -200,28 +205,22 @@ async def _wait_for_connection(self):
             logger.debug("Waiting for connection...")
             await asyncio.sleep(1)
 
-    async def _ratelimit(self, unit, max_sleep_time = 120):
-        rate_limit = self.rate_limit.get(unit)
-        if not rate_limit:
-            self.rate_limit[unit] = rate_limit = {
-                "backoff": 0,
-                "last_attempt_time": 0
-            }
+    async def _ratelimit(self, unit: str, max_sleep_time: int = 120) -> None:
+        limit: Optional[RateLimit] = self.rate_limit.get(unit)
+        if not limit:
+            self.rate_limit[unit] = limit = RateLimit()
 
-        if time.time() - rate_limit["last_attempt_time"] > max_sleep_time:
+        if time.time() - limit.last_attempt_time > max_sleep_time:
             # reset backoff if action lasted more than max_sleep_time
-            rate_limit["backoff"] = 0
+            limit.backoff = 0
         else:
             # increase backoff
-            rate_limit["backoff"] = (
-                min(rate_limit["backoff"] * 2, max_sleep_time)
-                if rate_limit["backoff"] > 0
-                else 1
+            limit.backoff = (
+                min(limit.backoff * 2, max_sleep_time) if limit.backoff > 0 else 1
             )
-        logger.debug(
-            "Sleeping for " + str(rate_limit["backoff"]) + " seconds before " + unit)
-        await asyncio.sleep( rate_limit["backoff"])
-        rate_limit["last_attempt_time"] = time.time()
+        logger.debug("Sleeping for " + str(limit.backoff) + " seconds before " + unit)
+        await asyncio.sleep(limit.backoff)
+        limit.last_attempt_time = int(time.time())
 
     async def _subscribe(self):
         """