diff --git a/voussoirkit/ratelimiter.py b/voussoirkit/ratelimiter.py index ee5de1f..3926f0d 100644 --- a/voussoirkit/ratelimiter.py +++ b/voussoirkit/ratelimiter.py @@ -1,6 +1,11 @@ import threading import time +from voussoirkit import sentinel + +REJECT = sentinel.Sentinel('reject mode') +SLEEP = sentinel.Sentinel('sleep mode') + class Ratelimiter: ''' The Ratelimiter class is used to limit how often you perform some other @@ -19,7 +24,7 @@ class Ratelimiter: self, allowance, *, - mode='sleep', + mode=SLEEP, operation_cost=1, period=1, starting_balance=None, @@ -29,11 +34,11 @@ class Ratelimiter: Our spending balance per `period` seconds. mode: - 'sleep': + SLEEP: If we do not have the balance for an operation, sleep until we do. Then return True every time. - 'reject': + REJECT: If we do not have the balance for an operation, do nothing and return False. Otherwise subtract the cost and return True. @@ -65,19 +70,25 @@ class Ratelimiter: return x - if mode not in ('sleep', 'reject'): - raise ValueError(f'Invalid mode {repr(mode)}.') - self.allowance = positive(allowance, 'allowance') self.period = positive(period, 'period') self.operation_cost = positive(operation_cost, 'operation_cost') - self.mode = mode if starting_balance is None: self.balance = operation_cost else: self.balance = positive(starting_balance, 'starting_balance') + mode = { + 'reject': REJECT, + 'sleep': SLEEP, + }.get(mode, mode) + + if mode is not SLEEP and mode is not REJECT: + raise ValueError(f'mode should be SLEEP or REJECT, not {repr(mode)}.') + + self.mode = mode + self.lock = threading.Lock() self.last_operation = time.monotonic() @@ -95,7 +106,7 @@ class Ratelimiter: self.balance = min(self.balance, self.allowance) self.last_operation = now - if self.mode == 'reject' and self.balance < cost: + if self.mode is REJECT and self.balance < cost: success = False sleep_needed = 0 return (success, sleep_needed)