Decorators api

def add_documents(
    *,
    index_name: str,
    connection_info: Client | ConnectionInfo,
    batch_size: int | None = None,
    primary_key: str | None = None,
    wait_for_task: bool = False,
    verify: bool = True,
) -> Callable:
    """Decorator that takes the returned documents from a function and adds them to Meilisearch.

    It is required that either an client or url is provided.

    Args:
        index_name: The name of the index to which the documents should be added.
        connection_info: Either an Client instance ConnectionInfo with information on how to
            connect to Meilisearch.
        batch_size: If provided the documents will be sent in batches of the specified size.
            Otherwise all documents are sent at once. Default = None.
        primary_key: The primary key of the documents. This will be ignored if already set.
            Defaults to None.
        wait_for_task: If set to `True` the decorator will wait for the document addition to finish
            indexing before returning, otherwise it will return right away. Default = False.
        verify: If set to `False` the decorator will not verify the SSL certificate of the server.

    Returns:
        The list of documents provided by the decorated function.

    Raises:
        MeilisearchCommunicationError: If there was an error communicating with the server.
        MeilisearchApiError: If the Meilisearch API returned an error.
        ValueError: If neither an async_client nor an url is provided.

    Examples:
        >>> from meilisearch_python_sdk import Client
        >>> from meilisearch_python_sdk.decorators import add_documents, ConnectionInfo
        >>>
        >>>
        >>> # With `Client`
        >>> client = Client(url="http://localhost:7700", api_key="masterKey")
        >>> @add_documents(index_name="movies", connection_info=client)
        >>> def my_function() -> list[dict[str, Any]]:
        >>>     return [{"id": 1, "title": "Test 1"}, {"id": 2, "title": "Test 2"}]
        >>>
        >>> # With `ConnectionInfo`
        >>> @add_documents(
                index_name="movies",
                connection_info=ConnectionInfo(url="http://localhost:7700", api_key="masterKey"),
            )
        >>> def my_function() -> list[dict[str, Any]]:
        >>>     return [{"id": 1, "title": "Test 1"}, {"id": 2, "title": "Test 2"}]
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
            result = func(*args, **kwargs)
            if isinstance(connection_info, Client):
                _add_documents(
                    connection_info,
                    index_name,
                    result,
                    batch_size,
                    primary_key,
                    wait_for_task,
                )
                return result

            decorator_client = Client(
                url=connection_info.url, api_key=connection_info.api_key, verify=verify
            )
            _add_documents(
                decorator_client,
                index_name,
                result,
                batch_size,
                primary_key,
                wait_for_task,
            )

            return result

        return wrapper

    return decorator

def async_add_documents(
    *,
    index_name: str,
    connection_info: AsyncClient | ConnectionInfo,
    batch_size: int | None = None,
    primary_key: str | None = None,
    wait_for_task: bool = False,
    verify: bool = True,
) -> Callable:
    """Decorator that takes the returned documents from a function and asynchronously adds them to Meilisearch.

    It is required that either an async_client or url is provided.

    Args:
        index_name: The name of the index to which the documents should be added.
        connection_info: Either an AsyncClient instance ConnectionInfo with information on how to
            connect to Meilisearch.
        batch_size: If provided the documents will be sent in batches of the specified size.
            Otherwise all documents are sent at once. Default = None.
        primary_key: The primary key of the documents. This will be ignored if already set.
            Defaults to None.
        wait_for_task: If set to `True` the decorator will wait for the document addition to finish
            indexing before returning, otherwise it will return right away. Default = False.
        verify: If set to `False` the decorator will not verify the SSL certificate of the server.

    Returns:
        The list of documents provided by the decorated function.

    Raises:
        MeilisearchCommunicationError: If there was an error communicating with the server.
        MeilisearchApiError: If the Meilisearch API returned an error.
        ValueError: If neither an async_client nor an url is provided.

    Examples:
        >>> from meilisearch_python_sdk import AsyncClient
        >>> from meilisearch_python_sdk.decorators import async_add_documents, ConnectionInfo
        >>>
        >>>
        >>> # with `AsyncClient`
        >>> client = AsyncClient(url="http://localhost:7700", api_key="masterKey")
        >>> @async_add_documents(index_name="movies", connection_info=client)
        >>> async def my_function() -> list[dict[str, Any]]:
        >>>     return [{"id": 1, "title": "Test 1"}, {"id": 2, "title": "Test 2"}]
        >>>
        >>> # with `ConnectionInfo`
        >>> @async_add_documents(
                index_name="movies",
                connection_info=ConnectionInfo(url="http://localhost:7700", api_key="masterKey"),
            )
        >>> async def my_function() -> list[dict[str, Any]]:
        >>>     return [{"id": 1, "title": "Test 1"}, {"id": 2, "title": "Test 2"}]
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
            result = await func(*args, **kwargs)
            if isinstance(connection_info, AsyncClient):
                await _async_add_documents(
                    connection_info,
                    index_name,
                    result,
                    batch_size,
                    primary_key,
                    wait_for_task,
                )
                return result

            async with AsyncClient(
                connection_info.url, connection_info.api_key, verify=verify
            ) as client:
                await _async_add_documents(
                    client, index_name, result, batch_size, primary_key, wait_for_task
                )

            return result

        return wrapper

    return decorator