@@ -57,6 +57,7 @@ async def can_do_action(
5757 rate_hz : Optional [float ] = None ,
5858 burst_count : Optional [int ] = None ,
5959 update : bool = True ,
60+ n_actions : int = 1 ,
6061 _time_now_s : Optional [int ] = None ,
6162 ) -> Tuple [bool , float ]:
6263 """Can the entity (e.g. user or IP address) perform the action?
@@ -76,6 +77,9 @@ async def can_do_action(
7677 burst_count: How many actions that can be performed before being limited.
7778 Overrides the value set during instantiation if set.
7879 update: Whether to count this check as performing the action
80+ n_actions: The number of times the user wants to do this action. If the user
81+ cannot do all of the actions, the user's action count is not incremented
82+ at all.
7983 _time_now_s: The current time. Optional, defaults to the current time according
8084 to self.clock. Only used by tests.
8185
@@ -124,17 +128,20 @@ async def can_do_action(
124128 time_delta = time_now_s - time_start
125129 performed_count = action_count - time_delta * rate_hz
126130 if performed_count < 0 :
127- # Allow, reset back to count 1
128- allowed = True
131+ performed_count = 0
129132 time_start = time_now_s
130- action_count = 1.0
131- elif performed_count > burst_count - 1.0 :
133+
134+ # This check would be easier read as performed_count + n_actions > burst_count,
135+ # but performed_count might be a very precise float (with lots of numbers
136+ # following the point) in which case Python might round it up when adding it to
137+ # n_actions. Writing it this way ensures it doesn't happen.
138+ if performed_count > burst_count - n_actions :
132139 # Deny, we have exceeded our burst count
133140 allowed = False
134141 else :
135142 # We haven't reached our limit yet
136143 allowed = True
137- action_count += 1.0
144+ action_count = performed_count + n_actions
138145
139146 if update :
140147 self .actions [key ] = (action_count , time_start , rate_hz )
@@ -182,6 +189,7 @@ async def ratelimit(
182189 rate_hz : Optional [float ] = None ,
183190 burst_count : Optional [int ] = None ,
184191 update : bool = True ,
192+ n_actions : int = 1 ,
185193 _time_now_s : Optional [int ] = None ,
186194 ):
187195 """Checks if an action can be performed. If not, raises a LimitExceededError
@@ -201,6 +209,9 @@ async def ratelimit(
201209 burst_count: How many actions that can be performed before being limited.
202210 Overrides the value set during instantiation if set.
203211 update: Whether to count this check as performing the action
212+ n_actions: The number of times the user wants to do this action. If the user
213+ cannot do all of the actions, the user's action count is not incremented
214+ at all.
204215 _time_now_s: The current time. Optional, defaults to the current time according
205216 to self.clock. Only used by tests.
206217
@@ -216,6 +227,7 @@ async def ratelimit(
216227 rate_hz = rate_hz ,
217228 burst_count = burst_count ,
218229 update = update ,
230+ n_actions = n_actions ,
219231 _time_now_s = time_now_s ,
220232 )
221233
0 commit comments