Repository: niconoe/pyinaturalist (python)

Issue:
Improve performance for auto-paginated functions
Functions that perform pagination like `get_all_observations()` could be improved. They should use `requests.Session` to take advantage of its connection pooling, and also request the largest possible page size.

Pull request:
Improved pagination
Closes #82 

Patch:
diff --git a/pyinaturalist/api_requests.py b/pyinaturalist/api_requests.py
index 7e406a5..45c239f 100644
--- a/pyinaturalist/api_requests.py
+++ b/pyinaturalist/api_requests.py
@@ -1,4 +1,5 @@
 """ Some common functions for HTTP requests used by both the Node and REST API modules """
+import threading
 from logging import getLogger
 from os import getenv
 from typing import Dict, List, Union
@@ -8,6 +9,7 @@ import requests
 
 import pyinaturalist
 from pyinaturalist.constants import WRITE_HTTP_METHODS
+from pyinaturalist.forge_utils import copy_signature
 from pyinaturalist.request_params import preprocess_request_params, validate_ids
 
 # Mock response content to return in dry-run mode
@@ -15,27 +17,18 @@ MOCK_RESPONSE = Mock(spec=requests.Response)
 MOCK_RESPONSE.json.return_value = {'results': [], 'total_results': 0}
 
 logger = getLogger(__name__)
+thread_local = threading.local()
 
 
-# TODO: Copy function signature of request(), add `url`, and apply to these 4 wrapper functions
-def delete(url: str, **kwargs) -> requests.Response:
-    """ Wrapper around :py:func:`requests.delete` that supports dry-run mode """
-    return request('DELETE', url, **kwargs)
-
-
-def get(url: str, **kwargs) -> requests.Response:
-    """ Wrapper around :py:func:`requests.get` that supports dry-run mode """
-    return request('GET', url, **kwargs)
-
-
-def post(url: str, **kwargs) -> requests.Response:
-    """ Wrapper around :py:func:`requests.post` that supports dry-run mode """
-    return request('POST', url, **kwargs)
-
-
-def put(url: str, **kwargs) -> requests.Response:
-    """ Wrapper around :py:func:`requests.put` that supports dry-run mode """
-    return request('PUT', url, **kwargs)
+def get_session() -> requests.Session:
+    """Get a Session object that will be reused across requests to take advantage of connection
+    pooling. This is especially relevant for large paginated requests. If used in a multi-threaded
+    context (for example, a :py:class:`~concurrent.futures.ThreadPoolExecutor`), a separate session
+    is used for each thread.
+    """
+    if not hasattr(thread_local, "session"):
+        thread_local.session = requests.Session()
+    return thread_local.session
 
 
 def request(
@@ -46,6 +39,7 @@ def request(
     ids: Union[str, List] = None,
     params: Dict = None,
     headers: Dict = None,
+    session: requests.Session = None,
     **kwargs,
 ) -> requests.Response:
     """Wrapper around :py:func:`requests.request` that supports dry-run mode and
@@ -59,11 +53,13 @@ def request(
         ids: One or more integer IDs used as REST resource(s) to request
         params: Requests parameters
         headers: Request headers
+        session: Existing Session object to use instead of creating a new one
 
     Returns:
         API response
     """
     # Set user agent and authentication headers, if specified
+    session = session or get_session()
     headers = headers or {}
     headers['Accept'] = 'application/json'
     headers['User-Agent'] = user_agent or pyinaturalist.user_agent
@@ -72,7 +68,7 @@ def request(
 
     params = preprocess_request_params(params)
 
-    # If one or more REST resources are requested, update the request URL accordignly
+    # If one or more REST resources are requested by ID, update the request URL accordingly
     if ids:
         url = url.rstrip('/') + '/' + validate_ids(ids)
 
@@ -82,7 +78,31 @@ def request(
         log_request(method, url, params=params, headers=headers, **kwargs)
         return MOCK_RESPONSE
     else:
-        return requests.request(method, url, params=params, headers=headers, **kwargs)
+        return session.request(method, url, params=params, headers=headers, **kwargs)
+
+
+@copy_signature(request, exclude='method')
+def delete(url: str, **kwargs) -> requests.Response:
+    """ Wrapper around :py:func:`requests.delete` that supports dry-run mode """
+    return request('DELETE', url, **kwargs)
+
+
+@copy_signature(request, exclude='method')
+def get(url: str, **kwargs) -> requests.Response:
+    """ Wrapper around :py:func:`requests.get` that supports dry-run mode """
+    return request('GET', url, **kwargs)
+
+
+@copy_signature(request, exclude='method')
+def post(url: str, **kwargs) -> requests.Response:
+    """ Wrapper around :py:func:`requests.post` that supports dry-run mode """
+    return request('POST', url, **kwargs)
+
+
+@copy_signature(request, exclude='method')
+def put(url: str, **kwargs) -> requests.Response:
+    """ Wrapper around :py:func:`requests.put` that supports dry-run mode """
+    return request('PUT', url, **kwargs)
 
 
 def is_dry_run_enabled(method: str) -> bool:
diff --git a/pyinaturalist/constants.py b/pyinaturalist/constants.py
index fb8fa16..89e0a9b 100644
--- a/pyinaturalist/constants.py
+++ b/pyinaturalist/constants.py
@@ -5,8 +5,8 @@ INAT_NODE_API_BASE_URL = 'https://api.inaturalist.org/v1/'
 INAT_BASE_URL = 'https://www.inaturalist.org'
 INAT_KEYRING_KEY = '/inaturalist'
 
-PER_PAGE_RESULTS = 30  # Number of records per page for paginated queries
-THROTTLING_DELAY = 1  # In seconds, support <1 floats such as 0.1
+PER_PAGE_RESULTS = 200  # Default number of records per page for paginated queries
+THROTTLING_DELAY = 1.0  # Delay between paginated queries, in seconds
 
 # Toggle dry-run mode: this will run and log mock HTTP requests instead of real ones
 DRY_RUN_ENABLED = False  # Mock all requests, including GET
@@ -25,4 +25,4 @@ RequestParams = Dict[str, Any]
 ResponseObject = Dict[str, Any]
 MultiInt = Union[int, List[int]]
 MultiStr = Union[str, List[str]]
-TemplateFunction = Any  # Cannot use Callable/Protocol, as they will not allow a mix of signatures
+TemplateFunction = Any  # Cannot use Callable/Protocol, as these will not allow a mix of signatures
diff --git a/pyinaturalist/forge_utils.py b/pyinaturalist/forge_utils.py
index 73fcb43..9d56be5 100644
--- a/pyinaturalist/forge_utils.py
+++ b/pyinaturalist/forge_utils.py
@@ -1,8 +1,8 @@
+"""Utilities built on top of ``python-forge`` that simplify defining API docs by combining
+function signature modification with docstring modification.
+This module makes ``python-forge`` optional; if not installed, these functions will quietly fail
+without modifying the target functions.
 """
-Utilities built on top of ``python-forge`` used to simplify usage for API docs, mainly by
-combining function signature modification with docstring modification.
-"""
-from functools import wraps
 from inspect import cleandoc
 from itertools import chain
 from logging import getLogger
@@ -13,7 +13,7 @@ from pyinaturalist.constants import TemplateFunction
 logger = getLogger(__name__)
 
 
-def document_request_params(template_functions: List[TemplateFunction]):
+def document_request_params(template_functions: List[TemplateFunction]) -> Callable:
     """Document a function with both docstrings and function signatures from one or more
     template functions.
 
@@ -50,7 +50,7 @@ def document_request_params(template_functions: List[TemplateFunction]):
     """
     template_functions = [*template_functions, _user_agent]
 
-    def f(func):
+    def wrapper(func):
         # Modify docstring
         func = copy_docstrings(func, template_functions)
 
@@ -60,14 +60,9 @@ def document_request_params(template_functions: List[TemplateFunction]):
         except ImportError:
             logger.debug('Forge not installed; skipping runtime API documentation')
 
-        # Call modified function
-        @wraps(func)
-        def g(*args, **kwargs):
-            return func(*args, **kwargs)
-
-        return g
+        return func
 
-    return f
+    return wrapper
 
 
 def copy_docstrings(
@@ -91,14 +86,14 @@ def copy_docstrings(
 
     # Insert 'Returns' section, if present
     if return_value_desc:
-        docstring += '\n\nReturns:\n    ' + return_value_desc
+        docstring += f'\n\nReturns:\n    {return_value_desc}'
 
     target_function.__doc__ = docstring
     return target_function
 
 
 def _split_docstring(docstring):
-    """ Split a docstring into a function description + return value description, if present. """
+    """Split a docstring into a function description + return value description, if present"""
     if 'Returns:' in docstring:
         function_desc, return_value_desc = docstring.split('Returns:')
     else:
@@ -108,6 +103,21 @@ def _split_docstring(docstring):
     return cleandoc(function_desc.strip()), cleandoc(return_value_desc.strip())
 
 
+def copy_signature(template_function: Callable, include=None, exclude=None) -> Callable:
+    """A wrapper around :py:func:`forge.copy` that silently fails if forge is not installed"""
+
+    def wrapper(target_function: Callable):
+        try:
+            import forge
+        except ImportError:
+            return target_function
+
+        revision = forge.copy(template_function, include=include, exclude=exclude)
+        return revision(target_function)
+
+    return wrapper
+
+
 def copy_signatures(
     target_function: Callable, template_functions: List[TemplateFunction]
 ) -> Callable:
@@ -122,7 +132,7 @@ def copy_signatures(
 
 
 def _get_combined_revision(template_functions: List[TemplateFunction]):
-    """ Create a :py:class:`forge.Revision` from the combined parameters of multiple functions """
+    """Create a :py:class:`forge.Revision` from the combined parameters of multiple functions"""
     import forge
 
     # Use forge.copy to create a revision for each template function
diff --git a/pyinaturalist/node_api.py b/pyinaturalist/node_api.py
index 453e492..65e2331 100644
--- a/pyinaturalist/node_api.py
+++ b/pyinaturalist/node_api.py
@@ -15,7 +15,6 @@ Functions
 from logging import getLogger
 from time import sleep
 from typing import List
-from urllib.parse import urljoin
 
 import requests
 
@@ -50,14 +49,14 @@ from pyinaturalist.response_format import (
 logger = getLogger(__name__)
 
 
-def make_inaturalist_api_get_call(endpoint: str, **kwargs) -> requests.Response:
+def node_api_get(endpoint: str, **kwargs) -> requests.Response:
     """Make an API call to iNaturalist.
 
     Args:
-        endpoint: The name of an endpoint not including the base URL e.g. 'observations'
+        endpoint: The name of an endpoint resource, not including the base URL e.g. 'observations'
         kwargs: Arguments for :py:func:`.api_requests.request`
     """
-    return get(urljoin(INAT_NODE_API_BASE_URL, endpoint), **kwargs)
+    return get(f'{INAT_NODE_API_BASE_URL}{endpoint}', **kwargs)
 
 
 # Controlled Terms
@@ -127,9 +126,7 @@ def get_controlled_terms(taxon_id: int = None, user_agent: str = None) -> JsonRe
     """
     # This is actually two endpoints, but they are so similar it seems best to combine them
     endpoint = 'controlled_terms/for_taxon' if taxon_id else 'controlled_terms'
-    response = make_inaturalist_api_get_call(
-        endpoint, params={'taxon_id': taxon_id}, user_agent=user_agent
-    )
+    response = node_api_get(endpoint, params={'taxon_id': taxon_id}, user_agent=user_agent)
 
     # controlled_terms/for_taxon returns a 422 if the specified taxon does not exist
     if response.status_code in (404, 422):
@@ -143,7 +140,7 @@ def get_controlled_terms(taxon_id: int = None, user_agent: str = None) -> JsonRe
 
 
 def get_observation(observation_id: int, user_agent: str = None) -> JsonResponse:
-    """Get details about an observation.
+    """Get details about a single observation by ID
 
     **API reference:** https://api.inaturalist.org/v1/docs/#!/Observations/get_observations_id
 
@@ -213,7 +210,7 @@ def get_observations(user_agent: str = None, **params) -> JsonResponse:
         JSON response containing observation records
     """
     validate_multiple_choice_param(params, 'order_by', NODE_OBS_ORDER_BY_PROPERTIES)
-    r = make_inaturalist_api_get_call('observations', params=params, user_agent=user_agent)
+    r = node_api_get('observations', params=params, user_agent=user_agent)
     r.raise_for_status()
 
     observations = r.json()
@@ -223,7 +220,8 @@ def get_observations(user_agent: str = None, **params) -> JsonResponse:
     return observations
 
 
-# TODO: Should use a requests Session here for several subsequent requests
+# TODO: Consolidate logic from get_all_*() functions into a generic pagination function
+#       This would have two variations: by page number (most common) and by ID (as seen below)
 @document_request_params([*docs._get_observations, docs._only_id])
 def get_all_observations(user_agent: str = None, **params) -> List[JsonResponse]:
     """Like :py:func:`get_observations()`, but handles pagination and returns all results in one
@@ -289,7 +287,7 @@ def get_observation_species_counts(user_agent: str = None, **params) -> JsonResp
     Returns:
         JSON response containing taxon records with counts
     """
-    r = make_inaturalist_api_get_call(
+    r = node_api_get(
         'observations/species_counts',
         params=params,
         user_agent=user_agent,
@@ -324,7 +322,7 @@ def get_all_observation_species_counts(user_agent: str = None, **params) -> List
     Returns:
         Combined list of taxon records with counts
     """
-    results = []  # type: List[JsonResponse]
+    results: List[JsonResponse] = []
     page = 1
 
     pagination_params = {
@@ -403,7 +401,7 @@ def get_places_by_id(place_id: MultiInt, user_agent: str = None) -> JsonResponse
     Returns:
         JSON response containing place records
     """
-    r = make_inaturalist_api_get_call('places', ids=place_id, user_agent=user_agent)
+    r = node_api_get('places', ids=place_id, user_agent=user_agent)
     r.raise_for_status()
 
     # Convert coordinates to floats
@@ -465,7 +463,7 @@ def get_places_nearby(user_agent: str = None, **params) -> JsonResponse:
     Returns:
         JSON response containing place records, divided into 'standard' and 'community' places.
     """
-    r = make_inaturalist_api_get_call('places/nearby', params=params, user_agent=user_agent)
+    r = node_api_get('places/nearby', params=params, user_agent=user_agent)
     r.raise_for_status()
     return convert_all_place_coordinates(r.json())
 
@@ -497,7 +495,7 @@ def get_places_autocomplete(q: str, user_agent: str = None) -> JsonResponse:
     Returns:
         JSON response containing place records
     """
-    r = make_inaturalist_api_get_call('places/autocomplete', params={'q': q}, user_agent=user_agent)
+    r = node_api_get('places/autocomplete', params={'q': q}, user_agent=user_agent)
     r.raise_for_status()
 
     # Convert coordinates to floats
@@ -549,7 +547,7 @@ def get_projects(user_agent: str = None, **params) -> JsonResponse:
         JSON response containing project records
     """
     validate_multiple_choice_param(params, 'order_by', PROJECT_ORDER_BY_PROPERTIES)
-    r = make_inaturalist_api_get_call('projects', params=params, user_agent=user_agent)
+    r = node_api_get('projects', params=params, user_agent=user_agent)
     r.raise_for_status()
 
     response = r.json()
@@ -583,7 +581,7 @@ def get_projects_by_id(
     Returns:
         JSON response containing project records
     """
-    r = make_inaturalist_api_get_call(
+    r = node_api_get(
         'projects',
         ids=project_id,
         params={'rule_details': rule_details},
@@ -623,7 +621,7 @@ def get_taxa(user_agent: str = None, **params) -> JsonResponse:
         JSON response containing taxon records
     """
     params = translate_rank_range(params)
-    r = make_inaturalist_api_get_call('taxa', params=params, user_agent=user_agent)
+    r = node_api_get('taxa', params=params, user_agent=user_agent)
     r.raise_for_status()
 
     taxa = r.json()
@@ -660,7 +658,7 @@ def get_taxa_by_id(taxon_id: MultiInt, user_agent: str = None) -> JsonResponse:
     Returns:
         JSON response containing taxon records
     """
-    r = make_inaturalist_api_get_call('taxa', ids=taxon_id, user_agent=user_agent)
+    r = node_api_get('taxa', ids=taxon_id, user_agent=user_agent)
     r.raise_for_status()
 
     taxa = r.json()
@@ -700,7 +698,7 @@ def get_taxa_autocomplete(user_agent: str = None, **params) -> JsonResponse:
         JSON response containing taxon records
     """
     params = translate_rank_range(params)
-    r = make_inaturalist_api_get_call('taxa/autocomplete', params=params, user_agent=user_agent)
+    r = node_api_get('taxa/autocomplete', params=params, user_agent=user_agent)
     r.raise_for_status()
     json_response = r.json()
 


Test patch:
diff --git a/test/test_api_requests.py b/test/test_api_requests.py
index f2f4db5..31dac1d 100644
--- a/test/test_api_requests.py
+++ b/test/test_api_requests.py
@@ -7,13 +7,18 @@ from pyinaturalist.api_requests import MOCK_RESPONSE, delete, get, post, put, re
 
 # Just test that the wrapper methods call requests.request with the appropriate HTTP method
 @pytest.mark.parametrize(
-    'function, http_method',
+    'http_func, http_method',
     [(delete, 'DELETE'), (get, 'GET'), (post, 'POST'), (put, 'PUT')],
 )
-@patch('pyinaturalist.api_requests.request')
-def test_http_methods(mock_request, function, http_method):
-    function('https://url', param='value')
-    mock_request.assert_called_with(http_method, 'https://url', param='value')
+@patch('pyinaturalist.api_requests.requests.Session.request')
+def test_http_methods(mock_request, http_func, http_method):
+    http_func('https://url', params={'key': 'value'})
+    mock_request.assert_called_with(
+        http_method,
+        'https://url',
+        params={'key': 'value'},
+        headers={'Accept': 'application/json', 'User-Agent': pyinaturalist.user_agent},
+    )
 
 
 # Test that the requests() wrapper passes along expected headers; just tests kwargs, not mock response
@@ -35,7 +40,7 @@ def test_http_methods(mock_request, function, http_method):
         ),
     ],
 )
-@patch('pyinaturalist.api_requests.requests.request')
+@patch('pyinaturalist.api_requests.requests.Session.request')
 def test_request_headers(mock_request, input_kwargs, expected_headers):
     request('GET', 'https://url', **input_kwargs)
     request_kwargs = mock_request.call_args[1]
@@ -75,9 +80,9 @@ def test_request_headers(mock_request, input_kwargs, expected_headers):
     ],
 )
 @patch('pyinaturalist.api_requests.getenv')
-@patch('pyinaturalist.api_requests.requests')
+@patch('pyinaturalist.api_requests.requests.Session.request')
 def test_request_dry_run(
-    mock_requests,
+    mock_request,
     mock_getenv,
     enabled_const,
     enabled_env,
@@ -102,11 +107,11 @@ def test_request_dry_run(
 
     # Verify that the request was or wasn't mocked based on settings
     if expected_real_request:
-        assert mock_requests.request.call_count == 1
-        assert response == mock_requests.request()
+        assert mock_request.call_count == 1
+        assert response == mock_request()
     else:
         assert response == MOCK_RESPONSE
-        assert mock_requests.request.call_count == 0
+        assert mock_request.call_count == 0
 
 
 # In addition to the test cases above, ensure that the request/response isn't altered with dry-run disabled
diff --git a/test/test_node_api.py b/test/test_node_api.py
index 43fceaf..19b2c8c 100644
--- a/test/test_node_api.py
+++ b/test/test_node_api.py
@@ -369,21 +369,21 @@ def test_get_taxa(requests_mock):
         ({'max_rank': 'species', 'rank': 'override_me'}, SPECIES_AND_LOWER),
     ],
 )
-@patch('pyinaturalist.node_api.make_inaturalist_api_get_call')
+@patch('pyinaturalist.node_api.get')
 def test_get_taxa_by_rank_range(
-    mock_inaturalist_api_get_call,
+    mock_get,
     params,
     expected_ranks,
 ):
     # Make sure custom rank params result in the correct 'rank' param value
     get_taxa(**params)
-    kwargs = mock_inaturalist_api_get_call.call_args[1]
+    kwargs = mock_get.call_args[1]
     requested_rank = kwargs['params']['rank']
     assert requested_rank == expected_ranks
 
 
 # This is just a spot test of a case in which boolean params should be converted
-@patch('pyinaturalist.api_requests.requests.request')
+@patch('pyinaturalist.api_requests.requests.Session.request')
 def test_get_taxa_by_name_and_is_active(request):
     get_taxa(q='Lixus bardanae', is_active=False)
     request_kwargs = request.call_args[1]
diff --git a/test/test_request_params.py b/test/test_request_params.py
index 3fa0492..a877788 100644
--- a/test/test_request_params.py
+++ b/test/test_request_params.py
@@ -136,7 +136,7 @@ def test_preprocess_request_params(mock_bool, mock_datetime, mock_list, mock_str
 @patch('pyinaturalist.node_api.convert_all_coordinates', side_effect=lambda x: x)
 @patch('pyinaturalist.node_api.convert_all_place_coordinates', side_effect=lambda x: x)
 @patch('pyinaturalist.api_requests.preprocess_request_params')
-@patch('pyinaturalist.api_requests.requests.request')
+@patch('pyinaturalist.api_requests.requests.Session.request')
 def test_all_node_requests_use_param_conversion(
     request,
     preprocess_request_params,
@@ -159,7 +159,7 @@ def test_all_node_requests_use_param_conversion(
 @patch('pyinaturalist.rest_api.convert_all_coordinates')
 @patch('pyinaturalist.rest_api.sleep')
 @patch('pyinaturalist.api_requests.preprocess_request_params')
-@patch('pyinaturalist.api_requests.requests.request')
+@patch('pyinaturalist.api_requests.requests.Session.request')
 def test_all_rest_requests_use_param_conversion(
     request, preprocess_request_params, sleep, convert_all_place_coordinates, http_function
 ):