_pagination.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. # Copyright 2022-present, the HuggingFace Inc. team.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Contains utilities to handle pagination on Huggingface Hub."""
  15. from collections.abc import Iterable
  16. import httpx
  17. from . import get_session, hf_raise_for_status, http_backoff, logging
  18. logger = logging.get_logger(__name__)
  19. def paginate(path: str, params: dict, headers: dict) -> Iterable:
  20. """Fetch a list of models/datasets/spaces and paginate through results.
  21. This is using the same "Link" header format as GitHub.
  22. See:
  23. - https://requests.readthedocs.io/en/latest/api/#requests.Response.links
  24. - https://docs.github.com/en/rest/guides/traversing-with-pagination#link-header
  25. """
  26. session = get_session()
  27. r = session.get(path, params=params, headers=headers)
  28. hf_raise_for_status(r)
  29. yield from r.json()
  30. # Follow pages
  31. # Next link already contains query params
  32. next_page = _get_next_page(r)
  33. while next_page is not None:
  34. logger.debug(f"Pagination detected. Requesting next page: {next_page}")
  35. r = http_backoff("GET", next_page, headers=headers)
  36. hf_raise_for_status(r)
  37. yield from r.json()
  38. next_page = _get_next_page(r)
  39. def _get_next_page(response: httpx.Response) -> str | None:
  40. return response.links.get("next", {}).get("url")