feat: Code Ocean Version 3.1 updates

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())

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"

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:
 

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()
+        )

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