CODEWITHSUNDEEP
pythondecorator
linkyndy
linkyndy

Reputation: 17920

Python decorator to limit number of calls

I would like to write a decorator that limits the number of calls to the wrapped function. Say, if I want the wrapped function to be called maximum 10 times, the decorator should execute that function the first 10 times, then it should return None instead.

Here's what I've come up with:

from functools import wraps


def max_calls(num):
     """Decorator which allows its wrapped function to be called `num` times"""

     def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            calls = getattr(wrapper, 'calls', 0)
            calls += 1

            if calls == num:
                return None

            setattr(wrapper, 'calls', calls)
            return func(*args, **kwargs)
        setattr(wrapper, 'calls', 0)
        return wrapper
    return decorator

However, this count calls properly, returns None when the limit is reached, but...it doesn't reset between program runs. That is, if I execute the program once, the counter reaches 5, and then re-execute the program, it continues from 5. What do I need to change so that the decorator works properly?

Upvotes: 2

Views: 3168

Answers (1)

Gareth Rees
Gareth Rees

Reputation: 65854

The problem is that you maintain just one set of call counts. But that means each Flask request shares call counts with all the other requests. What you need to do is to maintain a separate set of call counts for each Flask request.

From reading the API documentation it looks as though the way to make this work is to carry out these three steps:

  1. Make a subclass of flask.Request that can store your function call counts:

    import collections
import flask

class MyRequest(flask.Request):
    """A subclass of Request that maintains function call counts."""
    def __init__(self, *args, **kwargs):
        super(MyRequest, self).__init__(*args, **kwargs)
        self.call_counts = collections.defaultdict(int)

  2. Set request_class to your subclass when you initialize your application:

    app = flask.Flask(__name__, request_class=MyRequest, ...)

  3. Rewrite your decorator to store its counts in the global flask.request object:

    import functools

def max_calls(num, default=None):
    """Decorator which allows its wrapped function to be called at most `num`
    times per Flask request, but which then returns `default`.

    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if flask.request.call_counts[func] == num:
                return default
            flask.request.call_counts[func] += 1
            return func(*args, **kwargs)
        return wrapper
    return decorator

But having written that, it would be remiss for me not to point out that your question seems very strange. Why do you want to restrict the number of times a function can be called by a Flask request? Are you trying to do some kind of rate limiting? It seems likely that whatever you want to do can be done better using some other approach.

Upvotes: 1

Related Questions