view roundup/rate_limit.py @ 8543:1ffa1f42e1da

refactor: rework mime type comparison and clean code rest.py: accept application/* as match for application/json in non /binary_context rest path. allow defining default mime type to return when file/message is missing mime type. Make it a class variable to it can be changed from text/plain to text/markdown or whatever. extract code from determine_output_format() to create create_valid_content_types() method which returns a list of matching mime types for a given type/subtype. Eliminate mostly duplicate return statements by introducing a variable to specify valid mime types in error message. rest_common.py: Fix error messages that now return application/* as valid mime type. CHANGES.txt upgrading.txt rest.txt: top level notes and corrections. Also correct rst syntax on earlier change.
author John Rouillard <rouilj@ieee.org>
date Tue, 24 Mar 2026 21:30:47 -0400
parents 224ccb8b49ca
children
line wrap: on
line source

# Originaly from
# https://smarketshq.com/implementing-gcra-in-python-5df1f11aaa96?gi=4b9725f99bfa
# with imports, modifications for python 2, implementation of
# set/get_tat and marshaling as string, support for testonly
# and status method.

from datetime import datetime, timedelta

try:
    # used by python 3.11 and newer use tz aware dates
    from datetime import UTC
    dt_min = datetime.min.replace(tzinfo=UTC)
    # start of unix epoch
    dt_epoch = datetime(1970, 1, 1, tzinfo=UTC)
    fromisoformat = datetime.fromisoformat
except ImportError:
    # python 2.7 and older than 3.11 - use naive dates
    dt_min = datetime.min
    dt_epoch = datetime(1970, 1, 1)
    def fromisoformat(date):
        # only for naive dates
        return datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%f")

from roundup.anypy.datetime_ import utcnow


class RateLimit:  # pylint: disable=too-few-public-methods

    __slots__ = ("count", "period")

    def __init__(self, count, period):
        self.count = count
        self.period = period

    @property
    def inverse(self):
        return self.period.total_seconds() / self.count


class Gcra:

    __slots__ = ("memory",)

    def __init__(self):
        self.memory = {}

    def get_tat(self, key):
        # This should return a previous tat for the key or the current time.
        if key in self.memory:
            return self.memory[key]
        else:
            return dt_min

    def set_tat(self, key, tat):
        self.memory[key] = tat

    def get_tat_as_string(self, key):
        # get value as string:
        #  YYYY-MM-DDTHH:MM:SS.mmmmmm
        # to allow it to be marshalled/unmarshaled
        if key in self.memory:
            return self.memory[key].isoformat()
        else:
            return dt_min.isoformat()

    def set_tat_as_string(self, key, tat):
        # Take value as string and unmarshall:
        #  YYYY-MM-DDTHH:MM:SS.mmmmmm
        # to datetime
        self.memory[key] = fromisoformat(tat)

    def update(self, key, limit, testonly=False):
        '''Determine if the item associated with the key should be
           rejected given the RateLimit limit.
        '''
        now = utcnow()
        tat = max(self.get_tat(key), now)
        separation = (tat - now).total_seconds()
        max_interval = limit.period.total_seconds() - limit.inverse
        if separation > max_interval:
            reject = True
        else:
            reject = False
            if not testonly:
                new_tat = max(tat, now) + timedelta(seconds=limit.inverse)
                self.set_tat(key, new_tat)
        return reject

    def status(self, key, limit):
        '''Return status suitable for displaying as headers:
            X-RateLimit-Limit: calls allowed per period. Period/window
                is not specified in any api I found.
            X-RateLimit-Limit-Period: Non standard. Defines period in
                seconds for RateLimit-Limit.
            X-RateLimit-Remaining: How many calls are left in this window.
            X-RateLimit-Reset: window ends in this many seconds (not an
                 epoch timestamp) and all RateLimit-Limit calls are
                 available again.
            Retry-After: if user's request fails, this is the next time there
                 will be at least 1 available call to be consumed.
        '''

        ret = {}
        tat = self.get_tat(key)
        # static defined headers according to limit
        # all values are strings as that is required when used as headers
        ret['X-RateLimit-Limit'] = str(limit.count)
        ret['X-RateLimit-Limit-Period'] = str(
                                           int(
                                            limit.period.total_seconds())
                                          )

        # status of current limit as of now
        now = utcnow()

        current_count = int((limit.period - (tat - now)).total_seconds() /
                            limit.inverse)
        ret['X-RateLimit-Remaining'] = str(min(current_count, limit.count))

        # tat_in_epochsec = (tat - datetime(1970, 1, 1)).total_seconds()
        seconds_to_tat = (tat - now).total_seconds()
        ret['X-RateLimit-Reset'] = str(max(seconds_to_tat, 0))
        ret['X-RateLimit-Reset-date'] = "%s" % tat
        ret['Now'] = str((now - dt_epoch).total_seconds())
        ret['Now-date'] = "%s" % now

        if self.update(key, limit, testonly=True):
            # A new request would be rejected if it was processes.
            # The user has to wait until an item is dequeued.
            # One item is dequeued every limit.inverse seconds.
            ret['Retry-After'] = str(int(limit.inverse))
            ret['Retry-After-Timestamp'] = "%s" % \
                    (now + timedelta(seconds=limit.inverse))
        else:
            # if we are not rejected, the user can post another
            # attempt immediately.
            # Do we even need this header if not rejected?
            # RFC implies this is used with a 503 (or presumably
            # 429 which may postdate the rfc). So if no error, no header?
            # ret['Retry-After'] = '0'
            # ret['Retry-After-Timestamp'] = str(ret['Now-date'])
            pass

        return ret

Roundup Issue Tracker: http://roundup-tracker.org/