diff --git a/src/codeocean/capsule.py b/src/codeocean/capsule.py index 3d9b2d9..e8ff516 100644 --- a/src/codeocean/capsule.py +++ b/src/codeocean/capsule.py @@ -5,17 +5,21 @@ from typing import Optional from requests_toolbelt.sessions import BaseUrlSession +from codeocean.components import Ownership, 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" @dataclass_json @@ -42,13 +46,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[Ownership] = 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 +108,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/components.py b/src/codeocean/components.py index 712b800..de80325 100644 --- a/src/codeocean/components.py +++ b/src/codeocean/components.py @@ -69,3 +69,9 @@ class SearchFilter: values: Optional[list[str | float]] = None range: Optional[SearchFilterRange] = None exclude: Optional[bool] = None + + +class Ownership(StrEnum): + Private = "private" + Shared = "shared" + Created = "created" diff --git a/src/codeocean/computation.py b/src/codeocean/computation.py index 723834b..976e66e 100644 --- a/src/codeocean/computation.py +++ b/src/codeocean/computation.py @@ -7,6 +7,7 @@ from time import sleep, time from codeocean.enum import StrEnum +from codeocean.folder import Folder, DownloadFileURL class ComputationState(StrEnum): @@ -54,14 +55,15 @@ class Computation: id: str created: int name: str - state: ComputationState run_time: int + state: ComputationState cloud_workstation: Optional[bool] = None data_assets: Optional[list[InputDataAsset]] = None - end_status: Optional[ComputationEndStatus] = None - has_results: Optional[bool] = None parameters: Optional[list[Param]] = None processes: Optional[list[PipelineProcess]] = None + end_status: Optional[ComputationEndStatus] = None + exit_code: Optional[int] = None + has_results: Optional[bool] = None @dataclass_json @@ -99,33 +101,6 @@ class RunParams: processes: Optional[list[PipelineProcessParams]] = None -@dataclass_json -@dataclass(frozen=True) -class FolderItem: - name: str - path: str - type: str - size: Optional[int] = None - - -@dataclass_json -@dataclass(frozen=True) -class Folder: - items: list[FolderItem] - - -@dataclass_json -@dataclass(frozen=True) -class ListFolderParams: - path: str - - -@dataclass_json -@dataclass(frozen=True) -class DownloadFileURL: - url: str - - @dataclass class Computations: diff --git a/src/codeocean/data_asset.py b/src/codeocean/data_asset.py index 7b8e10c..5fe07a3 100644 --- a/src/codeocean/data_asset.py +++ b/src/codeocean/data_asset.py @@ -6,15 +6,17 @@ from time import sleep, time from typing import Optional -from codeocean.components import SortOrder, SearchFilter, Permissions +from codeocean.components import Ownership, SortOrder, SearchFilter, Permissions from codeocean.computation import PipelineProcess, Param from codeocean.enum import StrEnum +from codeocean.folder import Folder, DownloadFileURL class DataAssetType(StrEnum): Dataset = "dataset" Result = "result" Combined = "combined" + Model = "model" class DataAssetState(StrEnum): @@ -26,10 +28,10 @@ class DataAssetState(StrEnum): @dataclass_json @dataclass(frozen=True) class Provenance: - commit: str - run_script: str - docker_image: str - capsule: str + commit: Optional[str] = None + run_script: Optional[str] = None + docker_image: Optional[str] = None + capsule: Optional[str] = None data_assets: Optional[list[str]] = None @@ -78,16 +80,18 @@ class DataAsset: state: DataAssetState type: DataAssetType last_used: int - app_parameters: Optional[list[AppParameter]] = None - custom_metadata: Optional[dict] = None - description: Optional[str] = None - failure_reason: Optional[str] = None files: Optional[int] = None - provenance: Optional[Provenance] = None size: Optional[int] = None - source_bucket: Optional[SourceBucket] = None + description: Optional[str] = None tags: Optional[list[str]] = None + provenance: Optional[Provenance] = None + source_bucket: Optional[SourceBucket] = None + custom_metadata: Optional[dict] = None + app_parameters: Optional[list[AppParameter]] = None contained_data_assets: Optional[list[ContainedDataAsset]] = None + last_transferred: Optional[int] = None + transfer_error: Optional[str] = None + failure_reason: Optional[str] = None @dataclass_json @@ -185,12 +189,6 @@ class DataAssetSortBy(StrEnum): Size = "size" -class DataAssetOwnership(StrEnum): - Private = "private" - Shared = "shared" - Created = "created" - - class DataAssetSearchOrigin(StrEnum): Internal = "internal" External = "external" @@ -199,16 +197,17 @@ class DataAssetSearchOrigin(StrEnum): @dataclass_json @dataclass(frozen=True) class DataAssetSearchParams: - limit: int - offset: int - archived: bool - favorite: bool query: Optional[str] = None + next_token: Optional[str] = None + offset: Optional[int] = None + limit: Optional[int] = None sort_field: Optional[DataAssetSortBy] = None sort_order: Optional[SortOrder] = None type: Optional[DataAssetType] = None - ownership: Optional[DataAssetOwnership] = None + ownership: Optional[Ownership] = None origin: Optional[DataAssetSearchOrigin] = None + favorite: Optional[bool] = None + archived: Optional[bool] = None filters: Optional[list[SearchFilter]] = None @@ -217,6 +216,7 @@ class DataAssetSearchParams: class DataAssetSearchResults: has_more: bool results: list[DataAsset] + next_token: Optional[str] = None @dataclass_json @@ -227,6 +227,13 @@ class ContainedDataAsset: size: Optional[int] = None +@dataclass_json +@dataclass(frozen=True) +class TransferDataParams: + target: Target + force: Optional[bool] = None + + @dataclass class DataAssets: @@ -301,3 +308,26 @@ def search_data_assets(self, search_params: DataAssetSearchParams) -> DataAssetS res = self.client.post("data_assets/search", json=search_params.to_dict()) return DataAssetSearchResults.from_dict(res.json()) + + def list_data_asset_files(self, data_asset_id: str, path: str = "") -> Folder: + data = { + "path": path, + } + + res = self.client.post(f"data_assets/{data_asset_id}/files", json=data) + + return Folder.from_dict(res.json()) + + def get_data_asset_file_download_url(self, data_asset_id: str, path: str) -> DownloadFileURL: + res = self.client.get( + f"data_assets/{data_asset_id}/files/download_url", + params={"path": path}, + ) + + return DownloadFileURL.from_dict(res.json()) + + def transfer_data_asset(self, data_asset_id: str, transfer_params: TransferDataParams): + self.client.post( + f"data_assets/{data_asset_id}/transfer", + json=transfer_params.to_dict() + ) diff --git a/src/codeocean/folder.py b/src/codeocean/folder.py new file mode 100644 index 0000000..43cac89 --- /dev/null +++ b/src/codeocean/folder.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from dataclasses_json import dataclass_json +from dataclasses import dataclass +from typing import Optional + + +@dataclass_json +@dataclass(frozen=True) +class FolderItem: + name: str + path: str + type: str + size: Optional[int] = None + + +@dataclass_json +@dataclass(frozen=True) +class Folder: + items: list[FolderItem] + + +@dataclass_json +@dataclass(frozen=True) +class ListFolderParams: + path: str + + +@dataclass_json +@dataclass(frozen=True) +class DownloadFileURL: + url: str