superb-ai-onprem 0.1.6__py3-none-any.whl → 0.3.0__py3-none-any.whl

spb_onprem/exports/params/create_export.py

@@ -0,0 +1,85 @@
+from typing import Union
+
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+from spb_onprem.data.params import DataListFilter
+
+
+def create_export_params(
+    dataset_id: str,
+    location: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+    name: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+    data_filter: Union[
+        UndefinedType,
+        DataListFilter,
+        dict
+    ] = Undefined,
+    data_count: Union[
+        UndefinedType,
+        int
+    ] = Undefined,
+    frame_count: Union[
+        UndefinedType,
+        int
+    ] = Undefined,
+    annotation_count: Union[
+        UndefinedType,
+        int
+    ] = Undefined,
+    meta: Union[
+        UndefinedType,
+        dict
+    ] = Undefined,
+):
+    """Create parameters for export creation.
+    
+    Args:
+        dataset_id (str): The ID of the dataset to create the export for.
+        location (Optional[str]): The location where the export will be stored.
+        name (Optional[str]): The name of the export.
+        data_filter (Optional[DataListFilter | dict]): The search filter of the data.
+        data_count (Optional[int]): The number of data items to export.
+        frame_count (Optional[int]): The number of frames to export.
+        annotation_count (Optional[int]): The number of annotations to export.
+        meta (Optional[dict]): The meta information for the export.
+        
+    Returns:
+        dict: Parameters for export creation
+        
+    Raises:
+        BadParameterError: If required parameters are missing
+    """
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+
+    params = {
+        "dataset_id": dataset_id,
+    }
+
+    if location is not Undefined:
+        params["location"] = location
+    if name is not Undefined:
+        params["name"] = name
+    if data_filter is not Undefined and data_filter is not None:
+        # Handle both DataListFilter objects and plain dicts
+        if isinstance(data_filter, DataListFilter):
+            params["data_filter"] = data_filter.model_dump(by_alias=True, exclude_unset=True)
+        else:
+            # Assume it's a dict and use it directly
+            params["data_filter"] = data_filter
+    if data_count is not Undefined:
+        params["data_count"] = data_count
+    if frame_count is not Undefined:
+        params["frame_count"] = frame_count
+    if annotation_count is not Undefined:
+        params["annotation_count"] = annotation_count
+    if meta is not Undefined:
+        params["meta"] = meta
+
+    return params 

spb_onprem/exports/params/delete_export.py

@@ -0,0 +1,30 @@
+from spb_onprem.exceptions import BadParameterError
+
+
+def delete_export_params(
+    dataset_id: str,
+    export_id: str,
+):
+    """Create parameters for export deletion.
+    
+    Args:
+        export_id (str): The ID of the export to delete.
+        
+    Returns:
+        dict: Parameters for export deletion
+        
+    Raises:
+        BadParameterError: If export_id is missing
+    """
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+
+    if export_id is None:
+        raise BadParameterError("Export ID is required")
+
+    params = {
+        "dataset_id": dataset_id,
+        "export_id": export_id,
+    }
+
+    return params

spb_onprem/exports/params/export.py

@@ -0,0 +1,30 @@
+from spb_onprem.exceptions import BadParameterError
+
+
+def get_export_params(
+    dataset_id: str,
+    export_id: str,
+):
+    """Create parameters for getting an export.
+    
+    Args:
+        export_id (str): The ID of the export to get.
+        dataset_id (Optional[str]): The ID of the dataset to get the export for.
+        slice_id (Optional[str]): The ID of the slice to get the export for.
+        
+    Returns:
+        dict: Parameters for getting an export
+        
+    Raises:
+        BadParameterError: If export_id is missing
+    """
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+
+    if export_id is None:
+        raise BadParameterError("Export ID is required")
+
+    return {
+        "dataset_id": dataset_id,
+        "export_id": export_id
+    } 

spb_onprem/exports/params/exports.py

@@ -0,0 +1,74 @@
+from typing import Optional, List, Union
+from spb_onprem.base_model import CustomBaseModel, Field
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+
+
+class ExportFilterOptions(CustomBaseModel):
+    """Options for filtering exports.
+    
+    Attributes:
+        id_in: Filter exports by list of IDs
+        name_contains: Filter exports by name containing this string
+        name: Filter exports by exact name match
+        location_contains: Filter exports by location containing this string
+        location: Filter exports by exact location match
+    """
+    id_in: Optional[List[str]] = Field(None, alias="idIn")
+    name_contains: Optional[str] = Field(None, alias="nameContains")
+    name: Optional[str] = Field(None, alias="name")
+    location_contains: Optional[str] = Field(None, alias="locationContains")
+    location: Optional[str] = Field(None, alias="location")
+
+
+class ExportFilter(CustomBaseModel):
+    """Filter criteria for export queries.
+    
+    Attributes:
+        must_filter: Conditions that must be met
+        not_filter: Conditions that must not be met
+    """
+    must_filter: Optional[ExportFilterOptions] = Field(None, alias="must")
+    not_filter: Optional[ExportFilterOptions] = Field(None, alias="not")
+
+
+def get_exports_params(
+    dataset_id: str,
+    export_filter: Union[
+        UndefinedType,
+        ExportFilter
+    ] = Undefined,
+    cursor: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+    length: int = 10,
+):
+    """Create parameters for getting exports.
+    
+    Args:
+        dataset_id (str): The ID of the dataset to get exports for.
+        export_filter (Optional[ExportFilter]): The filter to apply to the exports.
+        cursor (Optional[str]): The cursor to use for pagination.
+        length (int): The number of exports to get.
+        
+    Returns:
+        dict: Parameters for getting exports
+        
+    Raises:
+        BadParameterError: If dataset_id is missing
+    """
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+
+    params = {
+        "dataset_id": dataset_id,
+        "length": length
+    }
+
+    if export_filter is not Undefined and export_filter is not None:
+        params["filter"] = export_filter.model_dump(by_alias=True, exclude_unset=True)
+    if cursor is not Undefined:
+        params["cursor"] = cursor
+
+    return params 

spb_onprem/exports/params/update_export.py

@@ -0,0 +1,98 @@
+from typing import Union
+from datetime import datetime
+
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+from spb_onprem.data.params import DataListFilter
+
+
+def update_export_params(
+    dataset_id: str,
+    export_id: str,
+    location: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+    name: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+    data_filter: Union[
+        UndefinedType,
+        DataListFilter,
+        dict
+    ] = Undefined,
+    data_count: Union[
+        UndefinedType,
+        int
+    ] = Undefined,
+    frame_count: Union[
+        UndefinedType,
+        int
+    ] = Undefined,
+    annotation_count: Union[
+        UndefinedType,
+        int
+    ] = Undefined,
+    meta: Union[
+        UndefinedType,
+        dict
+    ] = Undefined,
+    completed_at: Union[
+        UndefinedType,
+        datetime
+    ] = Undefined,
+):
+    """Create parameters for export update.
+    
+    Args:
+        dataset_id (str): The ID of the dataset to update the export for.
+        export_id (str): The ID of the export to update.
+        location (Optional[str]): The location where the export will be stored.
+        name (Optional[str]): The name of the export.
+        data_filter (Optional[DataListFilter | dict]): The search filter of the data.
+        data_count (Optional[int]): The number of data items to export.
+        frame_count (Optional[int]): The number of frames to export.
+        annotation_count (Optional[int]): The number of annotations to export.
+        meta (Optional[dict]): The meta information for the export.
+        completed_at (Optional[datetime]): The completed time of the export.
+        
+    Returns:
+        dict: Parameters for export update
+        
+    Raises:
+        BadParameterError: If required parameters are missing
+    """
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+    if export_id is None:
+        raise BadParameterError("Export ID is required")
+
+    params = {
+        "dataset_id": dataset_id,
+        "export_id": export_id,
+    }
+
+    if location is not Undefined:
+        params["location"] = location
+    if name is not Undefined:
+        params["name"] = name
+    if data_filter is not Undefined and data_filter is not None:
+        # Handle both DataListFilter objects and plain dicts
+        if isinstance(data_filter, DataListFilter):
+            params["data_filter"] = data_filter.model_dump(by_alias=True, exclude_unset=True)
+        else:
+            # Assume it's a dict and use it directly
+            params["data_filter"] = data_filter
+    if data_count is not Undefined:
+        params["data_count"] = data_count
+    if frame_count is not Undefined:
+        params["frame_count"] = frame_count
+    if annotation_count is not Undefined:
+        params["annotation_count"] = annotation_count
+    if meta is not Undefined:
+        params["meta"] = meta
+    if completed_at is not Undefined:
+        params["completed_at"] = completed_at
+
+    return params 

spb_onprem/exports/queries.py

@@ -0,0 +1,158 @@
+from spb_onprem.exports.params import (
+    get_exports_params,
+    get_export_params,
+    create_export_params,
+    delete_export_params,
+    update_export_params,
+)
+
+
+class Schemas:
+    """Schemas for exports queries
+    """
+    EXPORT = '''
+        id
+        datasetId
+        name
+        dataFilter
+        location
+        dataCount
+        annotationCount
+        frameCount
+        meta
+        createdAt
+        createdBy
+        updatedAt
+        updatedBy
+        completedAt
+    '''
+
+    EXPORT_PAGE = f'''
+        exports {{
+            {EXPORT}
+        }}
+        next
+        totalCount
+    '''
+
+
+class Queries:
+    '''
+    Queries for exports
+    '''
+
+    GET_EXPORTS = {
+        "name": "exports",
+        "query": f'''
+            query exports(
+                $dataset_id: ID!,
+                $filter: ExportFilter,
+                $cursor: String,
+                $length: Int
+                $orderBy: ExportOrderBy
+            ) {{
+                exports(
+                    datasetId: $dataset_id,
+                    filter: $filter,
+                    cursor: $cursor,
+                    length: $length,
+                    orderBy: $orderBy
+                ) {{
+                    {Schemas.EXPORT_PAGE}
+                }}
+            }}
+        ''',
+        "variables": get_exports_params
+    }
+
+    GET_EXPORT = {
+        "name": "export",
+        "query": f'''
+            query export(
+                $dataset_id: ID!,
+                $export_id: ID!
+            ) {{
+                export(datasetId: $dataset_id, id: $export_id) {{
+                    {Schemas.EXPORT}
+                }}
+            }}
+        ''',
+        "variables": get_export_params
+    }
+
+    CREATE_EXPORT = {
+        "name": "createExport",
+        "query": f'''
+            mutation createExport(
+                $dataset_id: ID!,
+                $location: String,
+                $name: String,
+                $data_filter: JSONObject,
+                $data_count: Int,
+                $frame_count: Int,
+                $annotation_count: Int,
+                $meta: JSONObject
+            ) {{
+                createExport(
+                    datasetId: $dataset_id,
+                    location: $location,
+                    name: $name,
+                    dataFilter: $data_filter,
+                    dataCount: $data_count,
+                    frameCount: $frame_count,
+                    annotationCount: $annotation_count,
+                    meta: $meta
+                ) {{
+                    {Schemas.EXPORT}
+                }}
+            }}
+        ''',
+        "variables": create_export_params
+    }
+    
+    UPDATE_EXPORT = {
+        "name": "updateExport",
+        "query": f'''
+            mutation updateExport(
+                $dataset_id: ID!,
+                $export_id: ID!,
+                $location: String,
+                $name: String,
+                $data_filter: JSONObject,
+                $data_count: Int,
+                $frame_count: Int,
+                $annotation_count: Int,
+                $meta: JSONObject,
+                $completed_at: DateTime
+            ) {{
+                updateExport(
+                    datasetId: $dataset_id,
+                    id: $export_id,
+                    location: $location,
+                    name: $name,
+                    dataFilter: $data_filter,
+                    dataCount: $data_count,
+                    frameCount: $frame_count,
+                    annotationCount: $annotation_count,
+                    meta: $meta,
+                    completedAt: $completed_at
+                ) {{    
+                    {Schemas.EXPORT}
+                }}
+            }}
+        ''',
+        "variables": update_export_params
+    }
+
+    DELETE_EXPORT = {
+        "name": "deleteExport",
+        "query": '''
+            mutation deleteExport(
+                $dataset_id: ID!,
+                $export_id: ID!
+            ) {
+                deleteExport(datasetId: $dataset_id, id: $export_id)
+            }
+        ''',
+        "variables": delete_export_params
+    }

spb_onprem/exports/service.py

@@ -0,0 +1,224 @@
+from typing import Optional, List, Union
+
+from spb_onprem.base_service import BaseService
+from spb_onprem.base_types import (
+    Undefined,
+    UndefinedType,
+)
+from spb_onprem.data.params import DataListFilter
+
+from .entities import Export
+from .params import (
+    ExportFilter,
+)
+from .queries import Queries
+
+
+class ExportService(BaseService):
+    """Service class for handling export-related operations."""
+    
+    def create_export(
+        self,
+        dataset_id: str,
+        name: Union[
+            UndefinedType,
+            str
+        ] = Undefined,
+        data_filter: Union[
+            UndefinedType,
+            DataListFilter,
+            dict
+        ] = Undefined,
+        data_count: Union[
+            UndefinedType,
+            int
+        ] = Undefined,
+        frame_count: Union[
+            UndefinedType,
+            int
+        ] = Undefined,
+        annotation_count: Union[
+            UndefinedType,
+            int
+        ] = Undefined,
+        meta: Union[
+            UndefinedType,
+            dict
+        ] = Undefined,
+    ) -> Export:
+        """Create an export.
+        
+        Args:
+            dataset_id (str): The ID of the dataset to create the export for.
+            name (Optional[str]): The name of the export.
+            data_filter (Optional[DataListFilter | dict]): The search filter of the data.
+            data_count (Optional[int]): The number of data items to export.
+            frame_count (Optional[int]): The number of frames to export.
+            annotation_count (Optional[int]): The number of annotations to export.
+            meta (Optional[dict]): The meta information for the export.
+        """
+        response = self.request_gql(
+            Queries.CREATE_EXPORT,
+            Queries.CREATE_EXPORT["variables"](
+                dataset_id=dataset_id,
+                name=name,
+                data_filter=data_filter,
+                data_count=data_count,
+                frame_count=frame_count,
+                annotation_count=annotation_count,
+                meta=meta,
+            )
+        )
+        return Export.model_validate(response)
+    
+    def get_exports(
+        self,
+        dataset_id: str,
+        export_filter: Optional[ExportFilter] = None,
+        cursor: Optional[str] = None,
+        length: int = 10
+    ) -> tuple[List[Export], Optional[str], int]:
+        """Get exports.
+        
+        Args:
+            dataset_id (str): The ID of the dataset to get exports for.
+            export_filter (Optional[ExportsFilterOptions]): The filter to apply to the exports.
+            cursor (Optional[str]): The cursor to use for pagination.
+            length (int): The number of exports to get.
+        
+        Returns:
+            tuple[List[Export], Optional[str], int]: A tuple containing the exports, the next cursor, and the total count of exports.
+        """
+        response = self.request_gql(
+            Queries.GET_EXPORTS,
+            Queries.GET_EXPORTS["variables"](
+                dataset_id=dataset_id,
+                export_filter=export_filter,
+                cursor=cursor,
+                length=length,
+            )
+        )
+        exports_dict = response.get("exports", [])
+        return (
+            [Export.model_validate(export_dict) for export_dict in exports_dict],
+            response.get("next"),
+            response.get("totalCount"),
+        )
+
+    def get_export(
+        self,
+        dataset_id: str,
+        export_id: str,
+    ) -> Export:
+        """Get an export.
+        
+        Args:
+            dataset_id (str): The ID of the dataset to get the export for.
+            export_id (str): The ID of the export to get.
+        
+        Returns:
+            Export: The export object.
+        """
+        response = self.request_gql(
+            Queries.GET_EXPORT,
+            Queries.GET_EXPORT["variables"](
+                dataset_id=dataset_id,
+                export_id=export_id,
+            )
+        )
+        return Export.model_validate(response)
+
+    def update_export(
+        self,
+        dataset_id: str,
+        export_id: str,
+        location: Union[
+            UndefinedType,
+            str
+        ] = Undefined,
+        name: Union[
+            UndefinedType,
+            str
+        ] = Undefined,
+        data_filter: Union[
+            UndefinedType,
+            DataListFilter,
+            dict
+        ] = Undefined,
+        data_count: Union[
+            UndefinedType,
+            int
+        ] = Undefined,
+        frame_count: Union[
+            UndefinedType,
+            int
+        ] = Undefined,
+        annotation_count: Union[
+            UndefinedType,
+            int
+        ] = Undefined,
+        meta: Union[
+            UndefinedType,
+            dict
+        ] = Undefined,
+        completed_at: Union[
+            UndefinedType,
+            str
+        ] = Undefined,
+    ) -> Export:
+        """Update an export.
+        
+        Args:
+            dataset_id (str): The ID of the dataset to update the export for.
+            export_id (str): The ID of the export to update.
+            location (Optional[str]): The location where the export will be stored.
+            name (Optional[str]): The name of the export.
+            data_filter (Optional[DataListFilter | dict]): The search filter of the data.
+            data_count (Optional[int]): The number of data items to export.
+            frame_count (Optional[int]): The number of frames to export.
+            annotation_count (Optional[int]): The number of annotations to export.
+            meta (Optional[dict]): The meta information for the export.
+            completed_at (Optional[str]): The completed time of the export.
+        
+        Returns:
+            Export: The updated export object.
+        """
+        response = self.request_gql(
+            Queries.UPDATE_EXPORT,
+            Queries.UPDATE_EXPORT["variables"](
+                dataset_id=dataset_id,
+                export_id=export_id,
+                location=location,
+                name=name,
+                data_filter=data_filter,
+                data_count=data_count,
+                frame_count=frame_count,
+                annotation_count=annotation_count,
+                meta=meta,
+                completed_at=completed_at,
+            )
+        )
+        return Export.model_validate(response)
+
+    def delete_export(
+        self,
+        dataset_id: str,
+        export_id: str,
+    ) -> bool:
+        """Delete an export.
+        
+        Args:
+            dataset_id (str): The ID of the dataset to delete the export for.
+            export_id (str): The ID of the export to delete.
+        
+        Returns:
+            bool: True if the export was deleted, False otherwise.
+        """
+        response = self.request_gql(
+            Queries.DELETE_EXPORT,
+            Queries.DELETE_EXPORT["variables"](
+                dataset_id=dataset_id,
+                export_id=export_id,
+            )
+        )
+        return response