mirror of
https://github.com/element-hq/synapse
synced 2024-08-16 19:20:52 +00:00
Use contextlib.contextmanager instead of a custom class
This commit is contained in:
parent
9d9b230501
commit
3077cb2915
1 changed files with 12 additions and 22 deletions
|
@ -20,6 +20,7 @@ from synapse.api.errors import LimitExceededError
|
||||||
from synapse.util.async import sleep
|
from synapse.util.async import sleep
|
||||||
|
|
||||||
import collections
|
import collections
|
||||||
|
import contextlib
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
|
|
||||||
|
@ -112,16 +113,19 @@ class _PerHostRatelimiter(object):
|
||||||
or self.request_times
|
or self.request_times
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@contextlib.contextmanager
|
||||||
def ratelimit(self):
|
def ratelimit(self):
|
||||||
|
# `contextlib.contextmanager` takes a generator and turns it into a
|
||||||
|
# context manager. The generator should only yield once with a value
|
||||||
|
# to be returned by manager.
|
||||||
|
# Exceptions will be reraised at the yield.
|
||||||
|
|
||||||
request_id = object()
|
request_id = object()
|
||||||
|
ret = self._on_enter(request_id)
|
||||||
def on_enter():
|
try:
|
||||||
return self._on_enter(request_id)
|
yield ret
|
||||||
|
finally:
|
||||||
def on_exit(exc_type, exc_val, exc_tb):
|
self._on_exit(request_id)
|
||||||
return self._on_exit(request_id)
|
|
||||||
|
|
||||||
return ContextManagerFunction(on_enter, on_exit)
|
|
||||||
|
|
||||||
def _on_enter(self, request_id):
|
def _on_enter(self, request_id):
|
||||||
time_now = self.clock.time_msec()
|
time_now = self.clock.time_msec()
|
||||||
|
@ -210,17 +214,3 @@ class _PerHostRatelimiter(object):
|
||||||
deferred.callback(None)
|
deferred.callback(None)
|
||||||
except KeyError:
|
except KeyError:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
class ContextManagerFunction(object):
|
|
||||||
def __init__(self, on_enter, on_exit):
|
|
||||||
self.on_enter = on_enter
|
|
||||||
self.on_exit = on_exit
|
|
||||||
|
|
||||||
def __enter__(self):
|
|
||||||
if self.on_enter:
|
|
||||||
return self.on_enter()
|
|
||||||
|
|
||||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
||||||
if self.on_exit:
|
|
||||||
return self.on_exit(exc_type, exc_val, exc_tb)
|
|
||||||
|
|
Loading…
Reference in a new issue