diff --git a/src/codeocean/capsule.py b/src/codeocean/capsule.py index e8ff516..784ac34 100644 --- a/src/codeocean/capsule.py +++ b/src/codeocean/capsule.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from dataclasses_json import dataclass_json -from typing import Optional +from typing import Optional, Iterator from requests_toolbelt.sessions import BaseUrlSession from codeocean.components import Ownership, SortOrder, SearchFilter @@ -113,3 +113,18 @@ def search_capsules(self, search_params: CapsuleSearchParams) -> CapsuleSearchRe res = self.client.post("capsules/search", json=search_params.to_dict()) return CapsuleSearchResults.from_dict(res.json()) + + def search_capsules_iterator(self, search_params: CapsuleSearchParams) -> Iterator[Capsules]: + params = search_params.to_dict() + while True: + response = self.search_capsules( + search_params=CapsuleSearchParams(**params) + ) + + for result in response.results: + yield result + + if not response.has_more: + return + + params["next_token"] = response.next_token