diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 5fe07a3..46ed1df 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 @@ -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,