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

{dataroom_client-1.0.1.post62.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.post62.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.post62.dev0 → dataroom_client-1.0.1.post63.dev0}/dataroom_client/client.py

@@ -11,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
 
@@ -23,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:
@@ -37,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
         )
@@ -51,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}")
@@ -63,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,
@@ -88,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]]
@@ -105,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}'
 
 
@@ -114,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 = (
@@ -133,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 = {}
@@ -162,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:
@@ -184,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:
@@ -205,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():
@@ -226,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.")
@@ -239,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)
@@ -261,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,
@@ -287,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,
@@ -316,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}"
@@ -393,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,
@@ -419,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,
@@ -448,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}"
@@ -526,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,
@@ -552,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,
@@ -581,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.
 
@@ -589,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:
@@ -684,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,
@@ -713,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.'))
@@ -779,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({
@@ -806,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')
         
@@ -852,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:
@@ -898,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:
@@ -973,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:
@@ -1014,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. '
@@ -1034,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, '
@@ -1052,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",
@@ -1082,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,
@@ -1104,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,
@@ -1133,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,
         }
@@ -1255,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,
@@ -1279,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.'
@@ -1308,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",
@@ -1317,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.'
@@ -1331,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",
@@ -1341,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",
@@ -1353,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",
@@ -1376,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",
@@ -1388,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({
@@ -1398,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",
@@ -1413,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",
@@ -1434,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",
@@ -1457,7 +1765,7 @@ class AsyncRunner:
     _lock = threading.Lock() # To ensure thread-safe initialization
 
     @classmethod
-    def _initialize(cls):
+    def _initialize(cls) -> None:
         """Initializes the background event loop and thread if not already done."""
         with cls._lock:
             if cls._thread is not None:
@@ -1476,10 +1784,13 @@ class AsyncRunner:
             logger.debug("Initialized ClassAsyncRunner background thread")
 
     @classmethod
-    def run(cls, coro):
+    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()
@@ -1488,7 +1799,7 @@ class AsyncRunner:
         return future.result()
 
     @classmethod
-    def shutdown(cls):
+    def shutdown(cls) -> None:
         """
         Cleanly stops the shared event loop.
         This is registered with atexit and called automatically.
@@ -1512,10 +1823,11 @@ class DataRoomClientSync:
     The official client of the DataRoom API using synchronous method and requests.
     """
 
-    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_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 = (
@@ -1526,10 +1838,8 @@ class DataRoomClientSync:
             raise DataRoomError("DataRoom api_url is not set")
         self._async_client = DataRoomClient(api_key=self.api_key, api_url=self.api_url, timeout=timeout)
 
-    def __getattr__(self, name):
-        """
-        Dynamically create sync methods for all methods of the async client.
-        """
+    def __getattr__(self, name) -> Any:
+        # Dynamically create sync methods for all methods of the async client.
         attr = getattr(self._async_client, name)
 
         if not callable(attr):
@@ -1544,12 +1854,20 @@ class DataRoomClientSync:
 
         return sync_wrapper
 
-    def __dir__(self) -> List[str]:
+    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:
+        """
+        Download an image from a URL.
+
+        @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.post62.dev0 → dataroom_client-1.0.1.post63.dev0}/pyproject.toml

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