diff --git a/opencage/geocoder.py b/opencage/geocoder.py index 4622ce5..9a61a4c 100644 --- a/opencage/geocoder.py +++ b/opencage/geocoder.py @@ -8,6 +8,7 @@ import requests import backoff from .version import __version__ +from datetime import datetime try: import aiohttp @@ -130,6 +131,9 @@ class OpenCageGeocode: """ session = None + ratelimit_limit = None + ratelimit_remaining = None + ratelimit_reset = None def __init__(self, key, protocol='https', domain=DEFAULT_DOMAIN, sslcontext=None): """Constructor.""" @@ -262,6 +266,26 @@ def _opencage_request(self, params): else: response = requests.get(self.url, params=params, headers=self._opencage_headers('requests')) # pylint: disable=missing-timeout + headers_keys = { + 'ratelimit_limit': { + 'header': 'x-ratelimit-limit', + 'conversor': lambda v: int(v) if v else self.ratelimit_limit + }, + 'ratelimit_remaining': { + 'header': 'x-ratelimit-remaining', + 'conversor': lambda v: int(v) if v else self.ratelimit_remaining + }, + 'ratelimit_reset': { + 'header': 'x-ratelimit-reset', + 'conversor': lambda v: datetime.fromtimestamp(int(v)) if v else self.ratelimit_reset + } + } + + for prop_name, prop_data in headers_keys.items(): + header_value = response.headers.get(prop_data['header']) + conversor = prop_data['conversor'] + setattr(self, prop_name, conversor(header_value)) + try: response_json = response.json() except ValueError as excinfo: diff --git a/test/test_ratelimit_properties.py b/test/test_ratelimit_properties.py new file mode 100644 index 0000000..5c12de8 --- /dev/null +++ b/test/test_ratelimit_properties.py @@ -0,0 +1,45 @@ +# encoding: utf-8 +from datetime import datetime +from pathlib import Path +import os +import httpretty +from httpretty import httprettified +from opencage.geocoder import OpenCageGeocode + +# reduce maximum backoff retry time from 120s to 1s +os.environ['BACKOFF_MAX_TIME'] = '1' + +geocoder = OpenCageGeocode('abcde') + + +@httprettified +def test_rate_limit_properties_no_headers(): + httpretty.register_uri( + httpretty.GET, + geocoder.url, + body=Path('test/fixtures/uk_postcode.json').read_text(encoding="utf-8") + ) + _ = geocoder.geocode("EC1M 5RF") + + assert geocoder.ratelimit_limit is None + assert geocoder.ratelimit_remaining is None + assert geocoder.ratelimit_reset is None + + +@httprettified +def test_rate_limit_properties(): + httpretty.register_uri( + httpretty.GET, + geocoder.url, + body=Path('test/fixtures/uk_postcode.json').read_text(encoding="utf-8"), + adding_headers={ + 'X-RateLimit-Limit': '2500', + 'X-RateLimit-Remaining': '2487', + 'X-RateLimit-Reset': '1402185600' + } + ) + _ = geocoder.geocode("EC1M 5RF") + + assert geocoder.ratelimit_limit == 2500 + assert geocoder.ratelimit_remaining == 2487 + assert geocoder.ratelimit_reset == datetime.fromtimestamp(1402185600)