mirror of
https://github.com/home-assistant/core.git
synced 2025-04-28 19:27:51 +00:00
Fix throttle to work on instance-level
This commit is contained in:
parent
8a04e1f5f4
commit
47fc1deecb
@ -233,35 +233,42 @@ class Throttle(object):
|
|||||||
self.limit_no_throttle = limit_no_throttle
|
self.limit_no_throttle = limit_no_throttle
|
||||||
|
|
||||||
def __call__(self, method):
|
def __call__(self, method):
|
||||||
lock = threading.Lock()
|
|
||||||
|
|
||||||
if self.limit_no_throttle is not None:
|
if self.limit_no_throttle is not None:
|
||||||
method = Throttle(self.limit_no_throttle)(method)
|
method = Throttle(self.limit_no_throttle)(method)
|
||||||
|
|
||||||
|
# We want to be able to differentiate between function and method calls
|
||||||
|
# All methods have the classname in their qualname seperated by a '.'
|
||||||
|
# Functions have a '.' in their qualname if defined inline, but will
|
||||||
|
# be prefixed by '.<locals>.' so we strip that out.
|
||||||
|
is_func = '.' not in method.__qualname__.split('.<locals>.')[-1]
|
||||||
|
|
||||||
@wraps(method)
|
@wraps(method)
|
||||||
def wrapper(*args, **kwargs):
|
def wrapper(*args, **kwargs):
|
||||||
"""
|
"""
|
||||||
Wrapper that allows wrapped to be called only once per min_time.
|
Wrapper that allows wrapped to be called only once per min_time.
|
||||||
If we cannot acquire the lock, it is running so return None.
|
If we cannot acquire the lock, it is running so return None.
|
||||||
"""
|
"""
|
||||||
if not lock.acquire(False):
|
# pylint: disable=protected-access
|
||||||
return None
|
host = wrapper if is_func else args[0]
|
||||||
try:
|
if not hasattr(host, '_throttle_lock'):
|
||||||
last_call = wrapper.last_call
|
host._throttle_lock = threading.Lock()
|
||||||
|
|
||||||
|
if not host._throttle_lock.acquire(False):
|
||||||
|
return None
|
||||||
|
|
||||||
|
last_call = getattr(host, '_throttle_last_call', None)
|
||||||
# Check if method is never called or no_throttle is given
|
# Check if method is never called or no_throttle is given
|
||||||
force = not last_call or kwargs.pop('no_throttle', False)
|
force = not last_call or kwargs.pop('no_throttle', False)
|
||||||
|
|
||||||
|
try:
|
||||||
if force or utcnow() - last_call > self.min_time:
|
if force or utcnow() - last_call > self.min_time:
|
||||||
result = method(*args, **kwargs)
|
result = method(*args, **kwargs)
|
||||||
wrapper.last_call = utcnow()
|
host._throttle_last_call = utcnow()
|
||||||
return result
|
return result
|
||||||
else:
|
else:
|
||||||
return None
|
return None
|
||||||
finally:
|
finally:
|
||||||
lock.release()
|
host._throttle_lock.release()
|
||||||
|
|
||||||
wrapper.last_call = None
|
|
||||||
|
|
||||||
return wrapper
|
return wrapper
|
||||||
|
|
||||||
|
@ -218,3 +218,14 @@ class TestUtil(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertEqual(3, len(calls1))
|
self.assertEqual(3, len(calls1))
|
||||||
self.assertEqual(2, len(calls2))
|
self.assertEqual(2, len(calls2))
|
||||||
|
|
||||||
|
def test_throttle_per_instance(self):
|
||||||
|
""" Test that the throttle method is done per instance of a class. """
|
||||||
|
|
||||||
|
class Tester(object):
|
||||||
|
@util.Throttle(timedelta(seconds=1))
|
||||||
|
def hello(self):
|
||||||
|
return True
|
||||||
|
|
||||||
|
self.assertTrue(Tester().hello())
|
||||||
|
self.assertTrue(Tester().hello())
|
||||||
|
Loading…
x
Reference in New Issue
Block a user