From db4b0a50c0f2d42b5148da68cf101d8c30ee0b57 Mon Sep 17 00:00:00 2001 From: jake-valsamis Date: Thu, 20 Feb 2025 16:18:13 -0500 Subject: [PATCH 1/3] add iterator function for search capsule API --- src/codeocean/capsule.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/codeocean/capsule.py b/src/codeocean/capsule.py index e8ff516..78f64e1 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 From 31efe4738c3503f4a5decb66354c226691f76c3b Mon Sep 17 00:00:00 2001 From: jake-valsamis Date: Thu, 20 Feb 2025 16:18:59 -0500 Subject: [PATCH 2/3] fix whitespaces --- src/codeocean/capsule.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/codeocean/capsule.py b/src/codeocean/capsule.py index 78f64e1..bbbddab 100644 --- a/src/codeocean/capsule.py +++ b/src/codeocean/capsule.py @@ -113,7 +113,7 @@ 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: From 14a1ccbb0fb71c17d87d9f4b81ece926d5140ea9 Mon Sep 17 00:00:00 2001 From: jake-valsamis Date: Thu, 20 Feb 2025 16:19:37 -0500 Subject: [PATCH 3/3] fix more whitespaces --- src/codeocean/capsule.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/codeocean/capsule.py b/src/codeocean/capsule.py index bbbddab..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, Iterator +from typing import Optional, Iterator from requests_toolbelt.sessions import BaseUrlSession from codeocean.components import Ownership, SortOrder, SearchFilter