From 6044f7080484a59b36df34b8048db02cbcb30d2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krist=C3=B3f=20T=C3=B3th?= <kristof.toth@avatao.com>
Date: Mon, 30 Jul 2018 17:54:26 +0200
Subject: [PATCH] Refactor RateLimiter to allow proper subclassing

---
 lib/tfw/decorators/rate_limiter.py | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/lib/tfw/decorators/rate_limiter.py b/lib/tfw/decorators/rate_limiter.py
index d1972b8..2f21eac 100644
--- a/lib/tfw/decorators/rate_limiter.py
+++ b/lib/tfw/decorators/rate_limiter.py
@@ -1,25 +1,33 @@
 # Copyright (C) 2018 Avatao.com Innovative Learning Kft.
 # All Rights Reserved. See LICENSE file for details.
 
-from functools import wraps
+from functools import wraps, partial
 from time import time, sleep
 
 
 class RateLimiter:
-    def __init__(self, rate_per_second):
+    def __init__(self, rate_per_second, action=sleep):
         self.min_interval = 1 / float(rate_per_second)
+        self.action = action
+        self.fun = None
         self.last_call = time()
 
     def __call__(self, fun):
         @wraps(fun)
         def wrapper(*args, **kwargs):
-            self._limit_rate()
-            fun(*args, **kwargs)
+            self.fun = partial(fun, *args, **kwargs)
+            limit_seconds = self._limit_rate()
+            if limit_seconds:
+                self.action(limit_seconds)
+                return
+            self.fun()
         return wrapper
 
     def _limit_rate(self):
-        since_last_call = time() - self.last_call
-        to_next_call = self.min_interval - since_last_call
+        seconds_since_last_call = time() - self.last_call
+        seconds_to_next_call = self.min_interval - seconds_since_last_call
         self.last_call = time()
-        if to_next_call > 0:
-            sleep(to_next_call)
+
+        if seconds_to_next_call > 0:
+            return seconds_to_next_call
+        return 0