throttle.py

   1 import time
   2 from basic_commands import reply
   3
   4 throttles = {}
   5 count, start_time = range(2)
   6
   7 # Implements throttling of the wrapped function, allowing it to be executed only
   8 # once every seconds_between seconds, but not enforcing this restriction until
   9 # enforce_at usages occur within seconds_between seconds of each other.
  10
  11 # Usage:
  12 # @throttle("web")
  13 # def xkcd(...):
  14 def throttle(group, seconds_between=300, enforce_at=10, quiet=False):
  15   # This is the function that performs the wrapping of xkcd.
  16   def throttle_function(func):
  17     # And this is the function that becomes the wrapped xkcd.
  18     def throttled(*args, **kwargs):
  19       throttle_info = throttles.setdefault(group, [0, 0])
  20       current_time = time.time()
  21
  22       # If we're more than seconds_between seconds after the start of this
  23       # cluster, we drop old occurances and increase the start_time to match.
  24       if throttle_info[start_time] + seconds_between < current_time:
  25         drop_num = (current_time - throttle_info[start_time]) // seconds_between
  26         if drop_num >= throttle_info[count]:
  27           throttle_info[:] = [0,0]
  28         else:
  29           throttle_info[count] -= drop_num
  30           throttle_info[start_time] += seconds_between * drop_num
  31
  32       # If we haven't hit our enforcement limit, we allow it.
  33       if throttle_info[count] < enforce_at:
  34         throttle_info[count] += 1
  35         if throttle_info[start_time] == 0:
  36           throttle_info[start_time] = current_time
  37
  38         # Run the throttled function.
  39         return func(*args, **kwargs)
  40       else:
  41         # Attempt to report the error.
  42         if not quiet and type(args[0]) == str and type(args[1]) == str:
  43           reply(args[0], args[1], "I'm sorry, but my A.O. unit has become overloaded.  Wait a few minutes and ask me again.")
  44
  45     # Pass the docstring through.
  46     throttled.__doc__ = func.__doc__
  47     return throttled
  48   return throttle_function