dataroom-client 1.0.1.post49.dev0__tar.gz → 1.0.1.post63.dev0__tar.gz

{dataroom_client-1.0.1.post49.dev0 → dataroom_client-1.0.1.post63.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dataroom-client
-Version: 1.0.1.post49.dev0
+Version: 1.0.1.post63.dev0
 Summary: A python client to interface with the Dataroom backend API
 Author: Ales Kocjancic
 Author-email: hi@ales.io

{dataroom_client-1.0.1.post49.dev0 → dataroom_client-1.0.1.post63.dev0}/dataroom_client/client.py

@@ -1,4 +1,8 @@
 import asyncio
+import functools
+import inspect
+import threading
+import atexit
 from datetime import datetime
 import json as json_module
 import logging
@@ -7,7 +11,7 @@ import uuid
 from enum import Enum
 from io import BytesIO
 import mimetypes
-from typing import List, TypedDict, Optional
+from typing import AsyncIterable, TypedDict, Optional, Any
 from urllib.parse import urljoin
 import httpx
 
@@ -19,11 +23,11 @@ logger = logging.getLogger(__name__)
 class DataRoomError(Exception):
     """Base exception class for DataRoomClient errors"""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         self.response = kwargs.pop("response", None)
         super().__init__(*args, **kwargs)
 
-    def __str__(self):
+    def __str__(self) -> str:
         if self.response:
             return f"{super().__str__()}\n{self.response.status_code}\n{self.response.text}"
         else:
@@ -33,7 +37,15 @@ class DataRoomError(Exception):
 class DataRoomFile:
     """A wrapper for a file-like object that can be used with DataRoomClient"""
 
-    def __init__(self, bytes_io, content_type, path=None, extension=None):
+    def __init__(self, bytes_io, content_type, path=None, extension=None) -> None:
+        """
+        Initializes a DataRoomFile object.
+
+        @param bytes_io: A file-like object (e.g., BytesIO) containing the file data.
+        @param content_type: The MIME type of the file (e.g., 'image/jpeg').
+        @param path: Optional. The original path of the file.
+        @param extension: Optional. The file extension (e.g., '.jpg'). If not provided, it's inferred from content_type.
+        """
         extension = (
             mimetypes.guess_extension(content_type) or "" if extension is None else extension
         )
@@ -47,7 +59,13 @@ class DataRoomFile:
         self.path = path
 
     @classmethod
-    def from_path(cls, path: str):
+    def from_path(cls, path: str) -> "DataRoomFile":
+        """
+        Creates a DataRoomFile from a local file path.
+
+        @param path: The absolute or relative path to the local file.
+        @return: A DataRoomFile instance.
+        """
         content_type, encoding = mimetypes.guess_type(path)
         if not content_type:
             raise DataRoomError(f"Could not guess content type for file: {path}")
@@ -59,7 +77,14 @@ class DataRoomFile:
             )
 
     @classmethod
-    def from_bytesio(cls, bytes_io, extension):
+    def from_bytesio(cls, bytes_io, extension) -> "DataRoomFile":
+        """
+        Creates a DataRoomFile from a BytesIO object.
+
+        @param bytes_io: A BytesIO object containing the file data.
+        @param extension: The file extension (e.g., '.jpg').
+        @return: A DataRoomFile instance.
+        """
         assert extension is not None, "Please provide a file extension"
         return DataRoomFile(
             bytes_io=bytes_io,
@@ -84,7 +109,7 @@ class ImageUpdate(TypedDict, total=False):
     id: str  # noqa: A003
     source: Optional[str]
     attributes: Optional[dict]
-    tags: Optional[List[str]]
+    tags: Optional[list[str]]
     coca_embedding: Optional[str]
     related_images: Optional[dict[str, str]]
     datasets: Optional[list[str]]
@@ -101,7 +126,7 @@ class ImageCreate(TypedDict, total=False):
     datasets: Optional[list[str]]
 
 
-def arg_deprecation_msg(arg_name, msg=''):
+def arg_deprecation_msg(arg_name, msg='') -> str:
     return f'DEPRECATION WARNING: Argument "{arg_name}" is deprecated, and will be removed in the future. {msg}'
 
 
@@ -110,10 +135,11 @@ class DataRoomClient:
     The official client of the DataRoom API. See notebooks for usage examples.
     """
 
-    def __init__(self, api_key=None, api_url=None, timeout=120):
+    def __init__(self, api_key=None, api_url=None, timeout=120) -> None:
         """
         @param api_key: API key for DataRoom API
         @param api_url: URL of the DataRoom backend API
+        @param timeout: Timeout for the API requests
         """
         self.api_key = api_key or os.environ.get("DATAROOM_API_KEY")
         self.api_url = (
@@ -129,7 +155,7 @@ class DataRoomClient:
 
     async def _make_request(
         self, url, params=None, method="GET", json=None, files=None, headers=None,
-    ):
+    ) -> dict:
         absolute_url = urljoin(self.api_url, url)
         if headers is None:
             headers = {}
@@ -158,7 +184,7 @@ class DataRoomClient:
 
     async def _make_paginated_request(
         self, url, limit=1000, params=None, method="GET", json=None, headers=None,
-    ):
+    ) -> list[dict]:
         items = []
         next_url = url
         while next_url:
@@ -180,7 +206,7 @@ class DataRoomClient:
 
     async def _make_paginated_request_iter(
         self, url, limit=1000, params=None, method="GET", json=None, headers=None,
-    ):
+    ) -> AsyncIterable[dict]:
         next_url = url
         returned_items = 0
         while next_url:
@@ -201,11 +227,11 @@ class DataRoomClient:
                 break
 
     @staticmethod
-    def _dict_filter_none(d: dict):
+    def _dict_filter_none(d: dict) -> dict:
         return {k: v for k, v in d.items() if v is not None}
 
     @staticmethod
-    def _get_attributes_filter(attributes: dict | None):
+    def _get_attributes_filter(attributes: dict | None) -> str | None:
         if not attributes:
             return None
         for key, val in attributes.items():
@@ -222,7 +248,7 @@ class DataRoomClient:
         return attrs_str
 
     @staticmethod
-    def _validate_vector(vector: str):
+    def _validate_vector(vector: str) -> None:
         err_msg = "Argument vector must be a string representing a list of 768 floats."
         if not isinstance(vector, str) or not len(vector) > 0:
             raise DataRoomError(f"{err_msg} Not a string.")
@@ -235,6 +261,12 @@ class DataRoomClient:
 
     @classmethod
     async def download_image_from_url(cls, image_url: str) -> DataRoomFile:
+        """
+        Downloads an image from a URL and returns it as a DataRoomFile.
+
+        @param image_url: The URL of the image to download.
+        @return: A DataRoomFile instance containing the downloaded image.
+        """
         try:
             async with httpx.AsyncClient() as client:
                 response = await client.get(image_url)
@@ -257,11 +289,11 @@ class DataRoomClient:
         self,
         limit: int | None = 1000,
         page_size: int = None,
-        fields: List[str] = None,
-        include_fields: List[str] = None,
-        exclude_fields: List[str] = None,
+        fields: list[str] = None,
+        include_fields: list[str] = None,
+        exclude_fields: list[str] = None,
         all_fields: bool = False,
-        return_latents: List[str] = None,
+        return_latents: list[str] = None,
         cache_ttl: int = None,
         partitions_count: int = None,
         partition: int = None,
@@ -283,15 +315,15 @@ class DataRoomClient:
         aspect_ratio__lt: float = None,
         aspect_ratio__lte: float = None,
         source: str = None,
-        sources: List[str] = None,
-        sources__ne: List[str] = None,
+        sources: list[str] = None,
+        sources__ne: list[str] = None,
         attributes: dict = None,
         has_attributes: list = None,
         lacks_attributes: list = None,
-        has_latents: List[str] = None,
-        lacks_latents: List[str] = None,
-        has_masks: List[str] = None,
-        lacks_masks: List[str] = None,
+        has_latents: list[str] = None,
+        lacks_latents: list[str] = None,
+        has_masks: list[str] = None,
+        lacks_masks: list[str] = None,
         tags: list = None,
         tags__ne: list = None,
         tags__all: list = None,
@@ -312,7 +344,23 @@ class DataRoomClient:
         datasets__all: list = None,
         datasets__ne_all: list = None,
         datasets__empty: bool = None,
-    ):
+    ) -> list[dict]:
+        """
+        Retrieves a paginated list of images, with optional filtering and field selection.
+
+        @param limit: The maximum number of images to return.
+        @param page_size: The number of images to return per page.
+        @param fields: A list of fields to return for each image. This overrides the default fields.
+        @param include_fields: A list of fields to include in the response, in addition to `fields` or the default fields.
+        @param exclude_fields: A list of fields to exclude from the response.
+        @param all_fields: If True and `fields` is None, returns all available fields for each image.
+        @param return_latents: A list of latent types to return for each image.
+        @param cache_ttl: The time-to-live for caching of this request in seconds.
+        @param partitions_count: The total number of partitions to divide the data into.
+        @param partition: The specific partition number to retrieve.
+        @param ...: Various filter parameters to narrow down the image search.
+        @return: A list of image dictionaries.
+        """
         headers = {}
         if cache_ttl:
             headers["Cache-Control"] = f"max-age={cache_ttl}"
@@ -389,11 +437,11 @@ class DataRoomClient:
         self,
         limit: int | None = 1000,
         page_size: int = None,
-        fields: List[str] = None,
-        include_fields: List[str] = None,
-        exclude_fields: List[str] = None,
+        fields: list[str] = None,
+        include_fields: list[str] = None,
+        exclude_fields: list[str] = None,
         all_fields: bool = False,
-        return_latents: List[str] = None,
+        return_latents: list[str] = None,
         cache_ttl: int = None,
         partitions_count: int = None,
         partition: int = None,
@@ -415,15 +463,15 @@ class DataRoomClient:
         aspect_ratio__lt: float = None,
         aspect_ratio__lte: float = None,
         source: str = None,
-        sources: List[str] = None,
-        sources__ne: List[str] = None,
+        sources: list[str] = None,
+        sources__ne: list[str] = None,
         attributes: dict = None,
         has_attributes: list = None,
         lacks_attributes: list = None,
-        has_latents: List[str] = None,
-        lacks_latents: List[str] = None,
-        has_masks: List[str] = None,
-        lacks_masks: List[str] = None,
+        has_latents: list[str] = None,
+        lacks_latents: list[str] = None,
+        has_masks: list[str] = None,
+        lacks_masks: list[str] = None,
         tags: list = None,
         tags__ne: list = None,
         tags__all: list = None,
@@ -444,7 +492,25 @@ class DataRoomClient:
         datasets__all: list = None,
         datasets__ne_all: list = None,
         datasets__empty: bool = None,
-    ):
+    ) -> AsyncIterable[dict]:
+        """
+        Retrieves an iterator of images, with optional filtering and field selection.
+
+        This method is useful for processing a large number of images without loading them all into memory at once.
+
+        @param limit: The maximum number of images to return.
+        @param page_size: The number of images to return per page.
+        @param fields: A list of fields to return for each image. This overrides the default fields.
+        @param include_fields: A list of fields to include in the response, in addition to `fields` or the default fields.
+        @param exclude_fields: A list of fields to exclude from the response.
+        @param all_fields: If True and `fields` is None, returns all available fields for each image.
+        @param return_latents: A list of latent types to return for each image.
+        @param cache_ttl: The time-to-live for caching of this request in seconds.
+        @param partitions_count: The total number of partitions to divide the data into.
+        @param partition: The specific partition number to retrieve.
+        @param ...: Various filter parameters to narrow down the image search.
+        @yields: An image dictionary.
+        """
         headers = {}
         if cache_ttl:
             headers["Cache-Control"] = f"max-age={cache_ttl}"
@@ -522,11 +588,11 @@ class DataRoomClient:
         self,
         limit: int | None = 1000,
         page_size: int = None,
-        fields: List[str] = None,
-        include_fields: List[str] = None,
-        exclude_fields: List[str] = None,
+        fields: list[str] = None,
+        include_fields: list[str] = None,
+        exclude_fields: list[str] = None,
         all_fields: bool = False,
-        return_latents: List[str] = None,
+        return_latents: list[str] = None,
         cache_ttl: int = None,
         prefix_length: int = None,
         num_prefixes: int = None,
@@ -548,15 +614,15 @@ class DataRoomClient:
         aspect_ratio__lt: float = None,
         aspect_ratio__lte: float = None,
         source: str = None,
-        sources: List[str] = None,
-        sources__ne: List[str] = None,
+        sources: list[str] = None,
+        sources__ne: list[str] = None,
         attributes: dict = None,
         has_attributes: list = None,
         lacks_attributes: list = None,
-        has_latents: List[str] = None,
-        lacks_latents: List[str] = None,
-        has_masks: List[str] = None,
-        lacks_masks: List[str] = None,
+        has_latents: list[str] = None,
+        lacks_latents: list[str] = None,
+        has_masks: list[str] = None,
+        lacks_masks: list[str] = None,
         tags: list = None,
         tags__ne: list = None,
         tags__all: list = None,
@@ -577,7 +643,7 @@ class DataRoomClient:
         datasets__all: list = None,
         datasets__ne_all: list = None,
         datasets__empty: bool = None,
-    ):
+    ) -> list[dict]:
         """
         Get a list of random images.
 
@@ -585,6 +651,19 @@ class DataRoomClient:
         num_prefixes to adjust the randomness factor. In general, a smaller prefix_length will give you more samples,
         but less random and a higher num_prefixes will give you more samples, but slow down the query. The default
         values are prefix_length=5 and num_prefixes=100.
+
+        @param limit: The maximum number of images to return.
+        @param page_size: The number of images to return per page.
+        @param fields: A list of fields to return for each image. This overrides the default fields.
+        @param include_fields: A list of fields to include in the response, in addition to `fields` or the default fields.
+        @param exclude_fields: A list of fields to exclude from the response.
+        @param all_fields: If True and `fields` is None, returns all available fields for each image.
+        @param return_latents: A list of latent types to return for each image.
+        @param cache_ttl: The time-to-live for caching of this request in seconds.
+        @param partitions_count: The total number of partitions to divide the data into.
+        @param partition: The specific partition number to retrieve.
+        @param ...: Various filter parameters to narrow down the image search.
+        @return: A list of image dictionaries.
         """
         headers = {}
         if cache_ttl:
@@ -680,15 +759,15 @@ class DataRoomClient:
         aspect_ratio__lt: float = None,
         aspect_ratio__lte: float = None,
         source: str = None,
-        sources: List[str] = None,
-        sources__ne: List[str] = None,
+        sources: list[str] = None,
+        sources__ne: list[str] = None,
         attributes: dict = None,
         has_attributes: list = None,
         lacks_attributes: list = None,
-        has_latents: List[str] = None,
-        lacks_latents: List[str] = None,
-        has_masks: List[str] = None,
-        lacks_masks: List[str] = None,
+        has_latents: list[str] = None,
+        lacks_latents: list[str] = None,
+        has_masks: list[str] = None,
+        lacks_masks: list[str] = None,
         tags: list = None,
         tags__ne: list = None,
         tags__all: list = None,
@@ -709,8 +788,15 @@ class DataRoomClient:
         datasets__all: list = None,
         datasets__ne_all: list = None,
         datasets__empty: bool = None,
-    ):
+    ) -> int:
+        """
+        Returns the total count of images based on the provided filters.
 
+        @param partitions_count: The total number of partitions to divide the data into.
+        @param partition: The specific partition number to retrieve.
+        @param ...: Various filter parameters to narrow down the image count.
+        @return: The total number of images matching the filters.
+        """
         if source is not None:
             sources = [source]
             logger.warning(arg_deprecation_msg('source', 'Please use "sources" instead.'))
@@ -775,12 +861,23 @@ class DataRoomClient:
     async def get_image(
         self,
         image_id: str,
-        fields: List[str] = None,
-        include_fields: List[str] = None,
-        exclude_fields: List[str] = None,
+        fields: list[str] = None,
+        include_fields: list[str] = None,
+        exclude_fields: list[str] = None,
         all_fields: bool = False,
-        return_latents: List[str] = None,
-    ):
+        return_latents: list[str] = None,
+    ) -> dict:
+        """
+        Retrieves a single image by its ID.
+
+        @param image_id: The UUID of the image to retrieve.
+        @param fields: A list of fields to return for each image. This overrides the default fields.
+        @param include_fields: A list of fields to include in the response, in addition to `fields` or the default fields.
+        @param exclude_fields: A list of fields to exclude from the response.
+        @param all_fields: If True and `fields` is None, returns all available fields for each image.
+        @param return_latents: A list of latent types to return for the image.
+        @return: A dictionary representing the image.
+        """
         return await self._make_request(
             url=f"images/{image_id}/",
             params=self._dict_filter_none({
@@ -802,7 +899,26 @@ class DataRoomClient:
         tags: list[str] = None,
         related_images: dict[str, str] | None = None,
         datasets: list[str] = None,
-    ):
+    ) -> dict:
+        """
+        Creates a new image from a local file or a URL.
+
+        @param image_id: Optional. The UUID for the new image.
+        @param source: The source of the image (e.g. a project or website name).
+        @param image_file: A DataRoomFile object for a local image.
+        @param image_url: A URL for a remote image.
+        @param attributes: A dictionary of attributes to associate with the image.
+        @param tags: A list of tags to associate with the image.
+        @param related_images: A dictionary mapping relation names to image IDs. E.g.
+            `{
+                "img1": "im2",
+                "img2": "im2",
+                "another image": "im3",
+            }`.
+        @param datasets: A list of versioned dataset slugs identifying the datasets to add the image to. E.g.
+            `["my-dataset/1", "my-dataset/2", "another-dataset/1"]`.
+        @return: A dictionary representing the newly created image.
+        """
         if not image_file and not image_url:
             raise DataRoomError('Please provide either an "image_file" or "image_url" field')
         
@@ -848,8 +964,14 @@ class DataRoomClient:
 
     async def create_images(
         self,
-        images: List[ImageCreate],
-    ):
+        images: list[ImageCreate],
+    ) -> list[dict]:
+        """
+        Creates multiple images in a single bulk request.
+
+        @param images: A list of ImageCreate dictionaries, each defining an image to create.
+        @return: A list of dictionaries representing the newly created images.
+        """
         files = []
         for i, image in enumerate(images):
             if 'id' not in image:
@@ -894,19 +1016,36 @@ class DataRoomClient:
         image_id: str,
         source: str = None,
         attributes: dict = None,
-        latents: List[LatentType] = None,
+        latents: list[LatentType] = None,
         tags: list[str] = None,
         coca_embedding: str = None,
         related_images: dict[str, str] | None = None,
         datasets: list[str] = None,
-    ):
+    ) -> dict:
         """
-        Update the image:
+        Update the image.
+
          * overwrite tags
          * merge attributes
          * merge latents
          * merge related_images
          * merge datasets
+
+        @param image_id: The UUID of the image to update.
+        @param source: The source of the image (e.g. a project or website name).
+        @param attributes: A dictionary of attributes to associate with the image.
+        @param latents: A list of latent types to associate with the image.
+        @param tags: A list of tags to associate with the image.
+        @param coca_embedding: A string representing a list of 768 floats, e.g. `"[0.12345,1.23456,...]"`.
+        @param related_images: A dictionary mapping relation names to image IDs. E.g.
+            `{
+                "img1": "im2",
+                "img2": "im2",
+                "another image": "im3",
+            }`.
+        @param datasets: A list of versioned dataset slugs identifying the datasets to add the image to. E.g.
+            `["my-dataset/1", "my-dataset/2", "another-dataset/1"]`.
+        @return: A dictionary representing the updated image.
         """
 
         if coca_embedding:
@@ -969,15 +1108,19 @@ class DataRoomClient:
 
     async def update_images(
         self,
-        images: List[ImageUpdate],
-    ):
+        images: list[ImageUpdate],
+    ) -> list[dict]:
         """
-        Bulk update images:
+        Bulk update images.
+
          * overwrite tags
          * merge attributes
          * merge latents
          * merge related_images
          * merge datasets
+
+        @param images: A list of ImageUpdate dictionaries, each defining an image to update.
+        @return: A list of dictionaries representing the updated images.
         """
         for image in images:
             if 'id' not in image:
@@ -1010,9 +1153,15 @@ class DataRoomClient:
         self,
         image_id: str,
         attributes: dict,
-    ):
+    ) -> dict:
         """
+        DEPRECATED: Adds or updates attributes for a single image. Please use `update_image` instead.
+
         Update attributes of an image, merging them with the existing attributes.
+
+        @param image_id: The UUID of the image to update.
+        @param attributes: A dictionary of attributes to associate with the image.
+        @return: A dictionary representing the updated image.
         """
         logger.warning(
             'DEPRECATION WARNING: Method "add_image_attributes" is deprecated, and will be removed in the future. '
@@ -1030,9 +1179,14 @@ class DataRoomClient:
     async def add_image_attributes_in_bulk(
         self,
         ids_to_attributes: dict[str, dict],
-    ):
+    ) -> list[dict]:
         """
+        DEPRECATED: Adds or updates attributes for multiple images in bulk. Please use `update_images` instead.
+
         Update attributes of a list of images, merging them with the existing attributes.
+
+        @param ids_to_attributes: A dictionary mapping image IDs to dictionaries of attributes.
+        @return: A list of dictionaries representing the updated images.
         """
         logger.warning(
             'DEPRECATION WARNING: Method "add_image_attributes_in_bulk" is deprecated, '
@@ -1048,18 +1202,36 @@ class DataRoomClient:
             ],
         )
 
-    async def delete_image(self, image_id: str):
+    async def delete_image(self, image_id: str) -> dict:
+        """
+        Deletes a single image by its ID.
+
+        @param image_id: The UUID of the image to delete.
+        """
         return await self._make_request(
             url=f"images/{image_id}/",
             method="DELETE",
         )
 
-    async def get_image_audit_logs(self, image_id: str):
+    async def get_image_audit_logs(self, image_id: str) -> list[dict]:
+        """
+        Retrieves the audit logs for a single image.
+
+        @param image_id: The UUID of the image.
+        @return: A list of audit log entries.
+        """
         return await self._make_request(
             url=f"images/{image_id}/audit_logs/",
         )
 
-    async def get_image_similarity(self, image_id_1: str, image_id_2: str):
+    async def get_image_similarity(self, image_id_1: str, image_id_2: str) -> dict:
+        """
+        Calculates the similarity score between two images.
+
+        @param image_id_1: The UUID of the first image.
+        @param image_id_2: The UUID of the second image.
+        @return: A dictionary containing the similarity score.
+        """
         response = await self._make_request(
             url=f"images/{image_id_1}/similarity/",
             method="POST",
@@ -1078,11 +1250,11 @@ class DataRoomClient:
         image_text: str = None,
         # options
         number=5,
-        fields: List[str] = None,
-        include_fields: List[str] = None,
-        exclude_fields: List[str] = None,
+        fields: list[str] = None,
+        include_fields: list[str] = None,
+        exclude_fields: list[str] = None,
         all_fields: bool = False,
-        return_latents: List[str] = None,
+        return_latents: list[str] = None,
         # filters
         short_edge: int | None = None,
         short_edge__gt: int = None,
@@ -1100,15 +1272,15 @@ class DataRoomClient:
         aspect_ratio__gte: float = None,
         aspect_ratio__lt: float = None,
         aspect_ratio__lte: float = None,
-        sources: List[str] = None,
-        sources__ne: List[str] = None,
+        sources: list[str] = None,
+        sources__ne: list[str] = None,
         attributes: dict = None,
         has_attributes: list = None,
         lacks_attributes: list = None,
-        has_latents: List[str] = None,
-        lacks_latents: List[str] = None,
-        has_masks: List[str] = None,
-        lacks_masks: List[str] = None,
+        has_latents: list[str] = None,
+        lacks_latents: list[str] = None,
+        has_masks: list[str] = None,
+        lacks_masks: list[str] = None,
         tags: list = None,
         tags__ne: list = None,
         tags__all: list = None,
@@ -1129,7 +1301,26 @@ class DataRoomClient:
         datasets__all: list = None,
         datasets__ne_all: list = None,
         datasets__empty: bool = None,
-    ):
+    ) -> list[dict]:
+        """
+        Finds images similar to a given image, vector, or text query.
+
+        You must provide exactly one of `image_id`, `image_file`, `image_vector`, or `image_text`.
+
+        @param image_id: Find images similar to the image with this UUID.
+        @param image_file: Find images similar to this local image file.
+        @param image_vector: Find images similar to this image embedding vector formatted as
+            a string of 768 floats, e.g. `"[0.12345,1.23456,...]"`.
+        @param image_text: Find images similar to this text query.
+        @param number: The number of similar images to return.
+        @param fields: A list of fields to return for each image. This overrides the default fields.
+        @param include_fields: A list of fields to include in the response, in addition to `fields` or the default fields.
+        @param exclude_fields: A list of fields to exclude from the response.
+        @param all_fields: If True and `fields` is None, returns all available fields for each image.
+        @param return_latents: A list of latent types to return for each image.
+        @param ...: Various filter and field selection parameters.
+        @return: A list of similar image dictionaries.
+        """
         search_args = {
             'image_id': image_id, 'image_file': image_file, 'image_vector': image_vector, 'image_text': image_text,
         }
@@ -1251,12 +1442,23 @@ class DataRoomClient:
         self,
         image_id: str,
         # options
-        fields: List[str] = None,
-        include_fields: List[str] = None,
-        exclude_fields: List[str] = None,
+        fields: list[str] = None,
+        include_fields: list[str] = None,
+        exclude_fields: list[str] = None,
         all_fields: bool = False,
-        return_latents: List[str] = None,
-    ):
+        return_latents: list[str] = None,
+    ) -> list[dict]:
+        """
+        Retrieves images related to a specific image.
+
+        @param image_id: The UUID of the image to find related images for.
+        @param fields: A list of fields to return for each image. This overrides the default fields.
+        @param include_fields: A list of fields to include in the response, in addition to `fields` or the default fields.
+        @param exclude_fields: A list of fields to exclude from the response.
+        @param all_fields: If True and `fields` is None, returns all available fields for each image.
+        @param return_latents: A list of latent types to return for each image.
+        @return: A list of related image dictionaries.
+        """
         params = self._dict_filter_none({
             "fields": ",".join(fields) if fields else None,
             "include_fields": ",".join(include_fields) if include_fields else None,
@@ -1275,7 +1477,16 @@ class DataRoomClient:
         latent_file: DataRoomFile,
         latent_type: str,
         is_mask=None,
-    ):
+    ) -> dict:
+        """
+        DEPRECATED: Attaches a latent representation file to an image. Please use `update_image` instead.
+
+        @param image_id: The UUID of the image to update.
+        @param latent_file: A DataRoomFile object containing the latent data.
+        @param latent_type: A string identifying the type of latent.
+        @param is_mask: Deprecated parameter.
+        @return: A dictionary representing the updated image.
+        """
         logger.warning(
             'DEPRECATION WARNING: Method "set_image_latent" is deprecated, and will be removed in the future. '
             'Please use "update_image" instead.'
@@ -1304,7 +1515,14 @@ class DataRoomClient:
             files=files,
         )
 
-    async def delete_image_latent(self, image_id: str, latent_type: str):
+    async def delete_image_latent(self, image_id: str, latent_type: str) -> dict:
+        """
+        Deletes a latent representation from an image.
+
+        @param image_id: The UUID of the image to update.
+        @param latent_type: The type of the latent to delete.
+        @return: A dictionary representing the updated image.
+        """
         return await self._make_request(
             url=f"images/{image_id}/delete_latent/",
             method="POST",
@@ -1313,7 +1531,14 @@ class DataRoomClient:
             },
         )
 
-    async def set_image_coca_embedding(self, image_id: str, vector: str):
+    async def set_image_coca_embedding(self, image_id: str, vector: str) -> dict:
+        """
+        DEPRECATED: Sets the CoCa embedding vector for an image. Please use `update_image` instead.
+
+        @param image_id: The UUID of the image to update.
+        @param vector: A string representation of the 768-float embedding vector.
+        @return: A dictionary representing the updated image.
+        """
         logger.warning(
             'DEPRECATION WARNING: Method "set_image_coca_embedding" is deprecated, and will be removed in the future. '
             'Please use "update_image" instead.'
@@ -1327,7 +1552,14 @@ class DataRoomClient:
             },
         )
     
-    async def aggregate_images(self, field, type):
+    async def aggregate_images(self, field, type) -> dict:
+        """
+        Performs an aggregation operation on a specified field across all images.
+
+        @param field: The field to aggregate on (e.g., 'source', 'aspect_ratio').
+        @param type: The type of aggregation to perform (e.g., 'value_counts').
+        @return: The result of the aggregation.
+        """
         return await self._make_request(
             url="images/aggregate/",
             method="POST",
@@ -1337,7 +1569,14 @@ class DataRoomClient:
             },
         )
     
-    async def bucket_images(self, field, size):
+    async def bucket_images(self, field, size) -> list[dict]:
+        """
+        Groups images into buckets based on a specified field and bucket size.
+
+        @param field: The field to bucket on (e.g., 'date_created').
+        @param size: The size or interval for each bucket (e.g., 'day', 'month').
+        @return: A list of buckets with counts.
+        """
         return await self._make_request(
             url="images/bucket/",
             method="POST",
@@ -1349,18 +1588,37 @@ class DataRoomClient:
 
     # -------------------- Tag API methods --------------------
 
-    async def get_tags(self, limit: int = 1000):
+    async def get_tags(self, limit: int = 1000) -> list[dict]:
+        """
+        Retrieves a list of all tags.
+
+        @param limit: The maximum number of tags to return.
+        @return: A list of tag dictionaries.
+        """
         return await self._make_paginated_request(
             url=f"tags/",
             limit=limit,
         )
 
-    async def get_tag(self, tag_id: str):
+    async def get_tag(self, tag_id: str) -> dict:
+        """
+        Retrieves a single tag by its ID.
+
+        @param tag_id: The ID of the tag to retrieve.
+        @return: A dictionary representing the tag.
+        """
         return await self._make_request(
             url=f"tags/{tag_id}/",
         )
 
-    async def create_tag(self, name: str, description: str = None):
+    async def create_tag(self, name: str, description: str = None) -> dict:
+        """
+        Creates a new tag.
+
+        @param name: The name of the new tag.
+        @param description: An optional description for the tag.
+        @return: A dictionary representing the newly created tag.
+        """
         return await self._make_request(
             url="tags/",
             method="POST",
@@ -1372,7 +1630,14 @@ class DataRoomClient:
             ),
         )
 
-    async def tag_images(self, image_ids: List[str], tag_names: List[str]):
+    async def tag_images(self, image_ids: list[str], tag_names: list[str]) -> list[dict]:
+        """
+        Associates a list of tags with a list of images.
+
+        @param image_ids: A list of image UUIDs to tag.
+        @param tag_names: A list of tag names to apply to the images.
+        @return: A list of dictionaries representing the tagged images.
+        """
         return await self._make_request(
             url="tags/tag_images/",
             method="PUT",
@@ -1384,7 +1649,14 @@ class DataRoomClient:
 
     # -------------------- Dataset API methods --------------------
 
-    async def get_datasets(self, slug: str = None, limit: int = 1000):
+    async def get_datasets(self, slug: str = None, limit: int = 1000) -> list[dict]:
+        """
+        Retrieves a list of datasets, optionally filtered by slug.
+
+        @param slug: Optional. Filter datasets by a specific slug.
+        @param limit: The maximum number of datasets to return.
+        @return: A list of dataset dictionaries.
+        """
         return await self._make_paginated_request(
             url=f"datasets/",
             params=self._dict_filter_none({
@@ -1394,11 +1666,25 @@ class DataRoomClient:
         )
 
     async def get_dataset(self, slug_version: str):
+        """
+        Retrieves a single dataset version by its slug and version.
+
+        @param slug_version: The identifier for the dataset version (e.g., "my-dataset/1").
+        @return: A dictionary representing the dataset.
+        """
         return await self._make_request(
             url=f"datasets/{slug_version}/",
         )
 
-    async def create_dataset(self, name: str, slug: str, description: str = None):
+    async def create_dataset(self, name: str, slug: str, description: str = None) -> dict:
+        """
+        Creates a new dataset.
+
+        @param name: The display name of the dataset.
+        @param slug: The URL-friendly slug for the dataset. E.g. `"my-dataset"`.
+        @param description: An optional description for the dataset.
+        @return: A dictionary representing the newly created dataset.
+        """
         return await self._make_request(
             url=f"datasets/",
             method="POST",
@@ -1409,19 +1695,38 @@ class DataRoomClient:
             },
         )
     
-    async def freeze_dataset(self, slug_version: str):
+    async def freeze_dataset(self, slug_version: str) -> dict:
+        """
+        Freezes a dataset version, making it immutable.
+
+        @param slug_version: The identifier for the dataset version to freeze, e.g. "my-dataset/1".
+        @return: A dictionary representing the frozen dataset.
+        """
         return await self._make_request(
             url=f"datasets/{slug_version}/freeze/",
             method="POST",
         )
     
-    async def unfreeze_dataset(self, slug_version: str):
+    async def unfreeze_dataset(self, slug_version: str) -> dict:
+        """
+        Unfreezes a dataset version, making it mutable again.
+
+        @param slug_version: The identifier for the dataset version to unfreeze, e.g. "my-dataset/1".
+        @return: A dictionary representing the unfrozen dataset.
+        """
         return await self._make_request(
             url=f"datasets/{slug_version}/unfreeze/",
             method="POST",
         )
 
-    async def dataset_add_images(self, slug_version: str, image_ids: List[str]):
+    async def dataset_add_images(self, slug_version: str, image_ids: list[str]) -> dict:
+        """
+        Adds a list of images to a dataset version.
+
+        @param slug_version: The identifier for the dataset version, e.g. "my-dataset/1".
+        @param image_ids: A list of image UUIDs to add to the dataset.
+        @return: A dictionary representing the updated dataset.
+        """
         return await self._make_request(
             url=f"datasets/{slug_version}/images/",
             method="POST",
@@ -1430,7 +1735,14 @@ class DataRoomClient:
             },
         )
 
-    async def dataset_remove_images(self, slug_version: str, image_ids: List[str]):
+    async def dataset_remove_images(self, slug_version: str, image_ids: list[str]) -> dict:
+        """
+        Removes a list of images from a dataset version.
+
+        @param slug_version: The identifier for the dataset version, e.g. "my-dataset/1".
+        @param image_ids: A list of image UUIDs to remove from the dataset.
+        @return: A dictionary representing the updated dataset.
+        """
         return await self._make_request(
             url=f"datasets/{slug_version}/images/",
             method="DELETE",
@@ -1440,15 +1752,82 @@ class DataRoomClient:
         )
 
 
+
+class AsyncRunner:
+    """
+    Manages a single, shared event loop in a background thread
+    to run async functions from a synchronous context using classmethods.
+
+    The shutdown method is automatically registered to be called on exit.
+    """
+    _loop: asyncio.AbstractEventLoop | None = None
+    _thread: threading.Thread | None = None
+    _lock = threading.Lock() # To ensure thread-safe initialization
+
+    @classmethod
+    def _initialize(cls) -> None:
+        """Initializes the background event loop and thread if not already done."""
+        with cls._lock:
+            if cls._thread is not None:
+                return
+
+            cls._loop = asyncio.new_event_loop()
+            cls._thread = threading.Thread(
+                target=cls._loop.run_forever,
+                daemon=True,
+                name="ClassAsyncRunnerThread"
+            )
+            cls._thread.start()
+            # Register the shutdown method to be called when the program exits.
+            # This is done here to ensure it's only registered once.
+            atexit.register(cls.shutdown)
+            logger.debug("Initialized ClassAsyncRunner background thread")
+
+    @classmethod
+    def run(cls, coro) -> Any:
+        """
+        Runs a coroutine on the shared background event loop and returns the result.
+        Initializes the loop on the first call.
+
+        @param coro: The coroutine to run.
+        @return: The result of the coroutine.
+        """
+        if cls._thread is None:
+            cls._initialize()
+
+        future = asyncio.run_coroutine_threadsafe(coro, cls._loop)
+        return future.result()
+
+    @classmethod
+    def shutdown(cls) -> None:
+        """
+        Cleanly stops the shared event loop.
+        This is registered with atexit and called automatically.
+        """
+        # The check for cls._loop is important because atexit might call this
+        # even if the runner was never initialized.
+        if cls._loop and cls._loop.is_running():
+            logger.debug("Shutting down ClassAsyncRunner background thread...")
+            cls._loop.call_soon_threadsafe(cls._loop.stop)
+            # It's good practice to have a timeout on join
+            cls._thread.join(timeout=5) 
+            cls._loop.close()
+            logger.debug("ClassAsyncRunner has been shut down.")
+        
+        cls._loop = None
+        cls._thread = None
+
+
 class DataRoomClientSync:
     """
     The official client of the DataRoom API using synchronous method and requests.
     """
 
-    def __init__(self, api_key=None, api_url=None):
+    def __init__(self, api_key=None, api_url=None, timeout=120) -> None:
         """
-        @param api_key: API key for DataRoom API
+        @param api_key: API key for DataRoom API.
         @param api_url: URL of the DataRoom backend API
+        @param timeout: Timeout for the requests to the DataRoom backend API
         """
         self.api_key = api_key or os.environ.get("DATAROOM_API_KEY")
         self.api_url = (
@@ -1457,140 +1836,38 @@ class DataRoomClientSync:
         )
         if not self.api_url:
             raise DataRoomError("DataRoom api_url is not set")
-        self._async_client = DataRoomClient(api_key=api_key, api_url=api_url)
+        self._async_client = DataRoomClient(api_key=self.api_key, api_url=self.api_url, timeout=timeout)
 
-    # -------------------- Private methods --------------------
+    def __getattr__(self, name) -> Any:
+        # Dynamically create sync methods for all methods of the async client.
+        attr = getattr(self._async_client, name)
 
-    @classmethod
-    def _run_sync(cls, coro):
-        try:
-            # Check if there's an existing running event loop
-            loop = asyncio.get_running_loop()
-        except RuntimeError:
-            # No running event loop, create a new one
-            return asyncio.run(coro)
-        else:
-            # A running event loop exists, use run_until_complete
-            return loop.run_until_complete(coro)
+        if not callable(attr):
+            return attr
 
-    def _make_request(self, *args, **kwargs):
-        return self._run_sync(self._async_client._make_request(*args, **kwargs))
+        @functools.wraps(attr)
+        def sync_wrapper(*args, **kwargs):
+            result = attr(*args, **kwargs)
+            if inspect.isawaitable(result):
+                return AsyncRunner.run(result)
+            return result
 
-    def _make_paginated_request(self, *args, **kwargs):
-        return self._run_sync(
-            self._async_client._make_paginated_request(*args, **kwargs)
-        )
+        return sync_wrapper
 
-    # -------------------- Utils --------------------
+    def __dir__(self) -> list[str]:
+        """
+        Provide a list of attributes for introspection and autocompletion in tools like IPython.
+        """
+        # include all attributes from the async client and the sync client.
+        return sorted(list(set(super().__dir__()) | set(dir(self._async_client))))
 
     @classmethod
     def download_image_from_url(cls, *args, **kwargs) -> DataRoomFile:
-        return cls._run_sync(DataRoomClient.download_image_from_url(*args, **kwargs))
-
-    # -------------------- Image API methods --------------------
-
-    def get_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_images(*args, **kwargs))
-
-    def get_images_iter(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_images_iter(*args, **kwargs))
-    
-    def get_random_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_random_images(*args, **kwargs))
-
-    def count_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.count_images(*args, **kwargs))
-
-    def get_image(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_image(*args, **kwargs))
-
-    def create_image(self, *args, **kwargs):
-        return self._run_sync(self._async_client.create_image(*args, **kwargs))
-
-    def create_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.create_images(*args, **kwargs))
-
-    def delete_image(self, *args, **kwargs):
-        return self._run_sync(self._async_client.delete_image(*args, **kwargs))
-
-    def get_image_audit_logs(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_image_audit_logs(*args, **kwargs))
-
-    def get_image_similarity(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_image_similarity(*args, **kwargs))
-
-    def get_related_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_related_images(*args, **kwargs))
-
-    def get_similar_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_similar_images(*args, **kwargs))
-
-    def set_image_latent(self, *args, **kwargs):
-        return self._run_sync(self._async_client.set_image_latent(*args, **kwargs))
-
-    def delete_image_latent(self, *args, **kwargs):
-        return self._run_sync(self._async_client.delete_image_latent(*args, **kwargs))
-
-    def update_image(self,*args, **kwargs):
-        return self._run_sync(self._async_client.update_image(*args, **kwargs))
-
-    def update_images(self,*args, **kwargs):
-        return self._run_sync(self._async_client.update_images(*args, **kwargs))
-
-    def add_image_attributes(self, *args, **kwargs):
-        return self._run_sync(self._async_client.add_image_attributes(*args, **kwargs))
-
-    def add_image_attributes_in_bulk(self, *args, **kwargs):
-        return self._run_sync(self._async_client.add_image_attributes_in_bulk(*args, **kwargs))
-
-    def set_image_coca_embedding(self, *args, **kwargs):
-        return self._run_sync(self._async_client.set_image_coca_embedding(*args, **kwargs))
-    
-    def aggregate_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.aggregate_images(*args, **kwargs))
-
-    def bucket_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.bucket_images(*args, **kwargs))
-
-    # -------------------- Tag API methods --------------------
-
-    def create_tag(self, *args, **kwargs):
-        return self._run_sync(self._async_client.create_tag(*args, **kwargs))
-
-    def get_tag(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_tag(*args, **kwargs))
-
-    def get_tags(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_tags(*args, **kwargs))
-
-    def tag_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.tag_images(*args, **kwargs))
-
-    # -------------------- Dataset API methods --------------------
-
-    def get_datasets(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_datasets(*args, **kwargs))
-
-    def get_dataset(self, *args, **kwargs):
-        return self._run_sync(self._async_client.get_dataset(*args, **kwargs))
-
-    def create_dataset(self, *args, **kwargs):
-        return self._run_sync(self._async_client.create_dataset(*args, **kwargs))
-    
-    def freeze_dataset(self, *args, **kwargs):
-        return self._run_sync(self._async_client.freeze_dataset(*args, **kwargs))
-    
-    def unfreeze_dataset(self, *args, **kwargs):
-        return self._run_sync(self._async_client.unfreeze_dataset(*args, **kwargs))
-
-    def dataset_add_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.dataset_add_images(*args, **kwargs))
-
-    def dataset_remove_images(self, *args, **kwargs):
-        return self._run_sync(self._async_client.dataset_remove_images(*args, **kwargs))
-
+        """
+        Download an image from a URL.
 
-for method_name in dir(DataRoomClient):
-    if not method_name.startswith("_"):
-        if not hasattr(DataRoomClientSync, method_name):
-            logger.warning(f"Missing implementation: DataRoomClientSync.{method_name}")
+        @param image_url: The URL of the image to download.
+        @return: A DataRoomFile instance containing the downloaded image.
+        """
+        # Class methods are not covered by the automatic wrapping of async methods in __getattr__.
+        return AsyncRunner.run(DataRoomClient.download_image_from_url(*args, **kwargs))

{dataroom_client-1.0.1.post49.dev0 → dataroom_client-1.0.1.post63.dev0}/pyproject.toml

@@ -6,7 +6,7 @@ authors = [
 ]
 readme = "README.md"
 dynamic = []
-version = "1.0.1.post49.dev0"
+version = "1.0.1.post63.dev0"
 
 [tool.poetry]