label-studio-sdk 1.0.8__py3-none-any.whl → 1.0.11__py3-none-any.whl

label_studio_sdk/projects/exports/client_ext.py

@@ -0,0 +1,200 @@
+import json
+import time
+import asyncio
+import typing
+import pandas as pd
+from .client import ExportsClient, AsyncExportsClient
+from io import BytesIO
+from label_studio_sdk.versions.client import VersionsClient, AsyncVersionsClient
+from label_studio_sdk.core.api_error import ApiError
+
+
+class ExportTimeoutError(ApiError):
+    
+    def __init__(self, export_snapshot):
+        super().__init__(
+            status_code=500,
+            body=(
+                f"Export job timed out after {timeout} seconds: "
+                f"unable to retrieve export job {export_snapshot.id}. "
+                f"Current status: {export_snapshot.status}. "
+                f"Try manually checking the running job with "
+                f"`ls.projects.exports.get(project_id={project_id}, export_pk={export_snapshot.id})`."
+            )
+        )
+        
+class ExportFailedError(ApiError):
+    
+    def __init__(self, export_snapshot):
+        super().__init__(
+            status_code=500,
+            body=f"Export failed: {export_snapshot}"
+        )
+
+
+def _check_status(export_snapshot, converted_format_id, status):
+    if converted_format_id:
+        converted_format = next((c for c in export_snapshot.converted_formats if c.id == converted_format_id), None)
+        if converted_format and converted_format.status == status:
+            return True
+    else:
+        if export_snapshot.status == status:
+            return True
+    return False
+
+
+class ExportsClientExt(ExportsClient):
+    
+    def _bytestream_to_fileobj(self, bytestream: typing.Iterable[bytes] | bytes) -> typing.BinaryIO:
+        buffer = BytesIO()
+        if isinstance(bytestream, typing.Iterable): 
+            for chunk in bytestream:
+                buffer.write(chunk)
+        else:
+            buffer.write(bytestream)
+        buffer.seek(0)
+        return buffer
+    
+    def _bytestream_to_binary(self, bytestream: typing.Iterable[bytes]) -> bytes:
+        fileobj = self._bytestream_to_fileobj(bytestream)
+        return fileobj.getvalue()
+    
+    def _bytestream_to_json(self, bytestream: typing.Iterable[bytes]) -> dict:
+        fileobj = self._bytestream_to_fileobj(bytestream)
+        return json.load(fileobj)
+    
+    def _bytestream_to_pandas(self, bytestream: typing.Iterable[bytes]) -> pd.DataFrame:
+        fileobj = self._bytestream_to_fileobj(bytestream)
+        return pd.read_csv(fileobj)
+    
+    def _poll_export(self, project_id, export_snapshot, converted_format_id, timeout):
+        start_time = time.time()
+        while not _check_status(export_snapshot, None, 'completed'):
+            export_snapshot = self.get(project_id, export_pk=export_snapshot.id)
+            if _check_status(export_snapshot, None, 'failed'):
+                raise ExportFailedError(export_snapshot)
+            if time.time() - start_time > timeout:
+                raise ExportTimeoutError(export_snapshot)
+            time.sleep(1)
+    
+    def _get_bytestream(
+        self,
+        project_id: int,
+        export_type: str,
+        timeout: int = 60,
+        create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
+    ):
+        version = VersionsClient(client_wrapper=self._client_wrapper).get()
+        
+        if version.edition == "Enterprise":
+            # Enterprise edition exports are async, so we need to wait for the export job to complete
+            export_snapshot = self.create(project_id, **(create_kwargs or {}))
+            # Poll for base (JSON) export to complete
+            self._poll_export(project_id, export_snapshot, None, timeout)
+            # Convert to requested format if not JSON
+            if export_type != "JSON":
+                converted_proc = self.convert(project_id, export_pk=export_snapshot.id, export_type=export_type, **(convert_kwargs or {}))
+                self._poll_export(project_id, export_snapshot, converted_proc.converted_format, timeout)
+                
+            bytestream = self.download(project_id, export_pk=export_snapshot.id, export_type=export_type, request_options={'chunk_size': 1024}, **(download_kwargs or {}))
+        else:
+            # Community edition exports are sync, so we can download the file immediately
+            bytestream = self.download_sync(project_id, export_type=export_type, download_all_tasks=True, download_resources=True)
+        return bytestream
+    
+    def as_file(self, project_id: int, export_type: str = "JSON", timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = self._get_bytestream(project_id, export_type, timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return self._bytestream_to_fileobj(bytestream)
+    
+    def as_binary(self, project_id: int, export_type: str = "JSON", timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = self._get_bytestream(project_id, export_type, timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return self._bytestream_to_binary(bytestream)
+    
+    def as_json(self, project_id: int, timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = self._get_bytestream(project_id, "JSON", timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return self._bytestream_to_json(bytestream)
+    
+    def as_pandas(self, project_id: int, timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = self._get_bytestream(project_id, "CSV", timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return self._bytestream_to_pandas(bytestream)
+    
+class AsyncExportsClientExt(AsyncExportsClient):
+    
+    async def _bytestream_to_fileobj(self, bytestream: typing.AsyncGenerator[bytes, None] | bytes):
+        """Convert bytestream to file-like object"""
+        fileobj = BytesIO()
+        if isinstance(bytestream, typing.AsyncGenerator):
+            async for chunk in bytestream:
+                fileobj.write(chunk)
+        else:
+            fileobj.write(bytestream)
+        fileobj.seek(0)
+        return fileobj
+
+    async def _bytestream_to_binary(self, bytestream):
+        """Convert bytestream to binary data"""
+        fileobj = await self._bytestream_to_fileobj(bytestream)
+        return fileobj.getvalue()
+
+    async def _bytestream_to_json(self, bytestream):
+        """Convert bytestream to JSON object"""
+        fileobj = await self._bytestream_to_fileobj(bytestream)
+        return json.load(fileobj)
+
+    async def _bytestream_to_pandas(self, bytestream):
+        """Convert bytestream to pandas DataFrame"""
+        fileobj = await self._bytestream_to_fileobj(bytestream)
+        return pd.read_csv(fileobj)
+    
+    async def _poll_export(self, project_id, export_snapshot, converted_format_id, timeout):
+        start_time = time.time()
+        while not _check_status(export_snapshot, None, 'completed'):
+            export_snapshot = await self.get(project_id, export_pk=export_snapshot.id)
+            if _check_status(export_snapshot, None, 'failed'):
+                raise ExportFailedError(export_snapshot)
+            if time.time() - start_time > timeout:
+                raise ExportTimeoutError(export_snapshot)
+            await asyncio.sleep(1)
+
+    async def _get_bytestream(
+        self,
+        project_id: int,
+        export_type: str,
+        timeout: int = 60,
+        create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
+    ):
+        version = await AsyncVersionsClient(client_wrapper=self._client_wrapper).get()
+        if version.edition == "Enterprise":
+            # Enterprise edition exports are async, so we need to wait for the export job to complete
+            export_snapshot = await self.create(project_id, **(create_kwargs or {}))
+            # Poll for base (JSON) export to complete
+            await self._poll_export(project_id, export_snapshot, None, timeout)
+            # Convert to requested format if not JSON
+            if export_type != "JSON":
+                converted_proc = await self.convert(project_id, export_pk=export_snapshot.id, export_type=export_type, **(convert_kwargs or {}))
+                await self._poll_export(project_id, export_snapshot, converted_proc.converted_format, timeout)
+            
+            bytestream = self.download(project_id, export_pk=export_snapshot.id, export_type=export_type, request_options={'chunk_size': 1024}, **(download_kwargs or {}))
+        else:
+            bytestream = self.download_sync(project_id, export_type=export_type, download_all_tasks=True, download_resources=True)
+        return bytestream
+    
+    async def as_file(self, project_id: int, export_type: str = "JSON", timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = await self._get_bytestream(project_id, export_type, timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return await self._bytestream_to_fileobj(bytestream)
+    
+    async def as_binary(self, project_id: int, export_type: str = "JSON", timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = await self._get_bytestream(project_id, export_type, timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return await self._bytestream_to_binary(bytestream)
+    
+    async def as_json(self, project_id: int, timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = await self._get_bytestream(project_id, "JSON", timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return await self._bytestream_to_json(bytestream)
+    
+    async def as_pandas(self, project_id: int, timeout: int = 60, create_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, convert_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, download_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        bytestream = await self._get_bytestream(project_id, "CSV", timeout, create_kwargs, convert_kwargs, download_kwargs)
+        return await self._bytestream_to_pandas(bytestream)

label_studio_sdk/projects/exports/types/__init__.py

@@ -0,0 +1,6 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from .exports_convert_response import ExportsConvertResponse
+from .exports_list_formats_response_item import ExportsListFormatsResponseItem
+
+__all__ = ["ExportsConvertResponse", "ExportsListFormatsResponseItem"]

label_studio_sdk/projects/exports/types/exports_convert_response.py

@@ -0,0 +1,24 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ....core.pydantic_utilities import UniversalBaseModel
+import typing
+from ....types.export_format import ExportFormat
+import pydantic
+from ....core.pydantic_utilities import IS_PYDANTIC_V2
+
+
+class ExportsConvertResponse(UniversalBaseModel):
+    export_type: typing.Optional[ExportFormat] = None
+    converted_format: typing.Optional[int] = pydantic.Field(default=None)
+    """
+    ID of the converted format
+    """
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

label_studio_sdk/projects/exports/types/exports_list_formats_response_item.py

@@ -0,0 +1,44 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ....core.pydantic_utilities import UniversalBaseModel
+import typing
+from ....types.export_format import ExportFormat
+import pydantic
+from ....core.pydantic_utilities import IS_PYDANTIC_V2
+
+
+class ExportsListFormatsResponseItem(UniversalBaseModel):
+    name: typing.Optional[ExportFormat] = None
+    title: typing.Optional[str] = pydantic.Field(default=None)
+    """
+    Export format title
+    """
+
+    description: typing.Optional[str] = pydantic.Field(default=None)
+    """
+    Export format description
+    """
+
+    link: typing.Optional[str] = pydantic.Field(default=None)
+    """
+    Export format documentation link
+    """
+
+    tags: typing.Optional[typing.List[str]] = pydantic.Field(default=None)
+    """
+    Export format tags
+    """
+
+    disabled: typing.Optional[bool] = pydantic.Field(default=None)
+    """
+    If true, the export format is not supported by the project.
+    """
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

label_studio_sdk/projects/pauses/__init__.py

@@ -0,0 +1,2 @@
+# This file was auto-generated by Fern from our API Definition.
+