Skip to content

Instantly share code, notes, and snippets.

@yuitest
Created May 11, 2014 11:54
Show Gist options
  • Save yuitest/e5222108f5f3bfd1eb9c to your computer and use it in GitHub Desktop.
Save yuitest/e5222108f5f3bfd1eb9c to your computer and use it in GitHub Desktop.
API 制限をかけるクラス(デコレータつき。)
# coding: utf8
from __future__ import unicode_literals, print_function, division
from functools import wraps
import redis
import time
class APILimit(object):
class APILimitExceed(Exception):
pass
def __init__(self, redis, max_calls, in_seconds, key_prefix=b''):
self.max_calls = max_calls
self.in_seconds = in_seconds
self.redis = redis
self.key_prefix = key_prefix
def call(self, key):
r = self.redis
c = r.llen(key)
if c is not None and c > self.max_calls:
return False
if r.exists(key):
r.rpushx(key, b'')
else:
with r.pipeline() as p:
p.multi()
p.rpush(key, b'')
p.expire(key, self.in_seconds)
p.execute()
return True
@classmethod
def decorate(
cls, redis_, key=(lambda x, *args, **kwargs: x), *args, **kwargs
):
limiter = APILimit(redis_, *args, **kwargs)
def wrapper(f):
@wraps(f)
def wrapped(*args, **kwargs):
if limiter.call(key(*args, **kwargs)):
return f(*args, **kwargs)
else:
raise APILimit.APILimitExceed
return wrapped
return wrapper
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment