diff --git a/src/codeocean/capsule.py b/src/codeocean/capsule.py index 3d9b2d9..e58e1dd 100644 --- a/src/codeocean/capsule.py +++ b/src/codeocean/capsule.py @@ -5,17 +5,27 @@ from typing import Optional from requests_toolbelt.sessions import BaseUrlSession +from codeocean.components import SortOrder, SearchFilter from codeocean.computation import Computation from codeocean.data_asset import DataAssetAttachParams, DataAssetAttachResults from codeocean.enum import StrEnum class CapsuleStatus(StrEnum): - NonPublished = "non-published" - Submitted = "submitted" - Publishing = "publishing" - Published = "published" - Verified = "verified" + NonRelease = "non_release" + Release = "release" + + +class CapsuleSortBy(StrEnum): + Created = "created" + LastAccessed = "last_accessed" + Name = "name" + + +class CapsuleOwnership(StrEnum): + Private = "private" + Shared = "shared" + Created = "created" @dataclass_json @@ -42,13 +52,37 @@ class Capsule: cloned_from_url: Optional[str] = None description: Optional[str] = None field: Optional[str] = None - keywords: Optional[list[str]] = None + tags: Optional[list[str]] = None original_capsule: Optional[OriginalCapsuleInfo] = None - published_capsule: Optional[str] = None + release_capsule: Optional[str] = None submission: Optional[dict] = None versions: Optional[list[dict]] = None +@dataclass_json +@dataclass(frozen=True) +class CapsuleSearchParams: + query: Optional[str] = None + next_token: Optional[str] = None + offset: Optional[int] = None + limit: Optional[int] = None + sort_field: Optional[CapsuleSortBy] = None + sort_order: Optional[SortOrder] = None + ownership: Optional[CapsuleOwnership] = None + status: Optional[CapsuleStatus] = None + favorite: Optional[bool] = None + archived: Optional[bool] = None + filters: Optional[list[SearchFilter]] = None + + +@dataclass_json +@dataclass(frozen=True) +class CapsuleSearchResults: + has_more: bool + results: list[Capsule] + next_token: Optional[str] = None + + @dataclass class Capsules: @@ -80,3 +114,8 @@ def detach_data_assets(self, capsule_id: str, data_assets: list[str]): f"capsules/{capsule_id}/data_assets/", json=data_assets, ) + + def search_capsules(self, search_params: CapsuleSearchParams) -> CapsuleSearchResults: + res = self.client.post("capsules/search", json=search_params.to_dict()) + + return CapsuleSearchResults.from_dict(res.json()) diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 7b8e10c..09a3054 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -15,6 +15,7 @@ class DataAssetType(StrEnum): Dataset = "dataset" Result = "result" Combined = "combined" + Model = "model" class DataAssetState(StrEnum): @@ -199,11 +200,12 @@ class DataAssetSearchOrigin(StrEnum): @dataclass_json @dataclass(frozen=True) class DataAssetSearchParams: - limit: int - offset: int - archived: bool - favorite: bool + limit: Optional[int] = None + offset: Optional[int] = None + archived: Optional[bool] = None + favorite: Optional[bool] = None query: Optional[str] = None + next_token: Optional[str] = None sort_field: Optional[DataAssetSortBy] = None sort_order: Optional[SortOrder] = None type: Optional[DataAssetType] = None @@ -217,6 +219,7 @@ class DataAssetSearchParams: class DataAssetSearchResults: has_more: bool results: list[DataAsset] + next_token: Optional[str] = None @dataclass_json