From 4d15f618c8ff75fe9e7b1d1014182eca9bd2a0cb Mon Sep 17 00:00:00 2001 From: jtyoung84 <104453205+jtyoung84@users.noreply.github.com> Date: Mon, 6 Jan 2025 15:53:36 -0800 Subject: [PATCH 1/4] feat: adds paginated search method to data asset api --- src/codeocean/data_asset.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 5fe07a3..1b14dea 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from requests_toolbelt.sessions import BaseUrlSession from time import sleep, time -from typing import Optional +from typing import Optional, Iterator from codeocean.components import Ownership, SortOrder, SearchFilter, Permissions from codeocean.computation import PipelineProcess, Param @@ -331,3 +331,15 @@ def transfer_data_asset(self, data_asset_id: str, transfer_params: TransferDataP f"data_assets/{data_asset_id}/transfer", json=transfer_params.to_dict() ) + + def paginate_search(self, data_asset_search_params: DataAssetSearchParams) -> Iterator[list[DataAsset]]: + has_more = True + search_params = data_asset_search_params.to_dict() + while has_more: + response = self.search_data_assets( + search_params=DataAssetSearchParams(**search_params) + ) + next_token = response.next_token + has_more = response.has_more + search_params["next_token"] = next_token + yield response.results From b10dd60e33da00039102640aa601cf97cf4e916b Mon Sep 17 00:00:00 2001 From: jtyoung84 <104453205+jtyoung84@users.noreply.github.com> Date: Mon, 6 Jan 2025 16:24:20 -0800 Subject: [PATCH 2/4] feat: return iterator of data assets --- src/codeocean/data_asset.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 1b14dea..2cf70cf 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -332,7 +332,7 @@ def transfer_data_asset(self, data_asset_id: str, transfer_params: TransferDataP json=transfer_params.to_dict() ) - def paginate_search(self, data_asset_search_params: DataAssetSearchParams) -> Iterator[list[DataAsset]]: + def paginate_search(self, data_asset_search_params: DataAssetSearchParams) -> Iterator[DataAsset]: has_more = True search_params = data_asset_search_params.to_dict() while has_more: @@ -342,4 +342,5 @@ def paginate_search(self, data_asset_search_params: DataAssetSearchParams) -> It next_token = response.next_token has_more = response.has_more search_params["next_token"] = next_token - yield response.results + for result in response.results: + yield result From 896499d34c3983678f2b2c03f3202146c8488b31 Mon Sep 17 00:00:00 2001 From: jtyoung84 <104453205+jtyoung84@users.noreply.github.com> Date: Tue, 7 Jan 2025 09:38:54 -0800 Subject: [PATCH 3/4] Update src/codeocean/data_asset.py Co-authored-by: Zvika Gart --- src/codeocean/data_asset.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 2cf70cf..c146084 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -332,15 +332,17 @@ def transfer_data_asset(self, data_asset_id: str, transfer_params: TransferDataP json=transfer_params.to_dict() ) - def paginate_search(self, data_asset_search_params: DataAssetSearchParams) -> Iterator[DataAsset]: - has_more = True - search_params = data_asset_search_params.to_dict() - while has_more: + def search_data_assets_iterator(self, search_params: DataAssetSearchParams) -> Iterator[DataAsset]: + params = search_params.to_dict() + while True: response = self.search_data_assets( - search_params=DataAssetSearchParams(**search_params) + search_params=DataAssetSearchParams(**params) ) - next_token = response.next_token - has_more = response.has_more - search_params["next_token"] = next_token + for result in response.results: yield result + + if not response.has_more: + return + + params["next_token"] = response.next_token From 29eafe8f84588ac985c19afcebf3af7acd3c795b Mon Sep 17 00:00:00 2001 From: jtyoung84 <104453205+jtyoung84@users.noreply.github.com> Date: Tue, 7 Jan 2025 09:44:36 -0800 Subject: [PATCH 4/4] refactor: moves search_data_asset_iterator function --- src/codeocean/data_asset.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index c146084..46ed1df 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -309,6 +309,21 @@ def search_data_assets(self, search_params: DataAssetSearchParams) -> DataAssetS return DataAssetSearchResults.from_dict(res.json()) + def search_data_assets_iterator(self, search_params: DataAssetSearchParams) -> Iterator[DataAsset]: + params = search_params.to_dict() + while True: + response = self.search_data_assets( + search_params=DataAssetSearchParams(**params) + ) + + for result in response.results: + yield result + + if not response.has_more: + return + + params["next_token"] = response.next_token + def list_data_asset_files(self, data_asset_id: str, path: str = "") -> Folder: data = { "path": path, @@ -331,18 +346,3 @@ def transfer_data_asset(self, data_asset_id: str, transfer_params: TransferDataP f"data_assets/{data_asset_id}/transfer", json=transfer_params.to_dict() ) - - def search_data_assets_iterator(self, search_params: DataAssetSearchParams) -> Iterator[DataAsset]: - params = search_params.to_dict() - while True: - response = self.search_data_assets( - search_params=DataAssetSearchParams(**params) - ) - - for result in response.results: - yield result - - if not response.has_more: - return - - params["next_token"] = response.next_token