content: #  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Pagination utilities."""
  
  from typing import Any, Callable, Generator, List, TypeVar
  
  from zenml.models import BaseIdentifiedResponse, Page
  
  AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
  
  
  def depaginate(
      list_method: Callable[..., Page[AnyResponse]], **kwargs: Any
  ) -> List[AnyResponse]:
      """Depaginate the results from a client or store method that returns pages.
  
      Args:
          list_method: The list method to depaginate.
          **kwargs: Arguments for the list method.
  
      Returns:
          A list of the corresponding Response Models.
      """
      page = list_method(**kwargs)
      items = list(page.items)
      while page.index < page.total_pages:
          kwargs["page"] = page.index + 1
          page = list_method(**kwargs)
          items += list(page.items)
  
      return items
  
  
  def depaginate_stream(
      list_method: Callable[..., Page[AnyResponse]], **kwargs: Any
  ) -> Generator[AnyResponse, None, None]:
      """Depaginate the results from a client or store method that returns pages.
  
      Args:
          list_method: The list method to depaginate.
          **kwargs: Arguments for the list method.
  
      Yields:
          A generator of the corresponding Response Models.
      """
      page = list_method(**kwargs)
  
      for item in page.items:
          yield item
  
      while page.index < page.total_pages:
          kwargs["page"] = page.index + 1
          page = list_method(**kwargs)
  
          for item in page.items:
              yield item
  