prefect-client 3.1.14__py3-none-any.whl → 3.2.0__py3-none-any.whl

prefect/artifacts.py

@@ -2,22 +2,27 @@
 Interface for creating and reading artifacts.
 """
 
-import asyncio
-import json  # noqa: I001
+from __future__ import annotations
+
+import json
 import math
 import warnings
-from typing import TYPE_CHECKING, Any, Optional, Union
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
 from uuid import UUID
 
 from typing_extensions import Self
 
+from prefect._internal.compatibility.async_dispatch import async_dispatch
+from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas.actions import ArtifactCreate as ArtifactRequest
 from prefect.client.schemas.actions import ArtifactUpdate
 from prefect.client.schemas.filters import ArtifactFilter, ArtifactFilterKey
+from prefect.client.schemas.objects import Artifact as ArtifactResponse
 from prefect.client.schemas.sorting import ArtifactSort
-from prefect.client.utilities import get_or_create_client
+from prefect.context import MissingContextError, get_run_context
 from prefect.logging.loggers import get_logger
-from prefect.utilities.asyncutils import sync_compatible
+from prefect.utilities.asyncutils import asyncnullcontext
 from prefect.utilities.context import get_task_and_flow_run_ids
 
 if TYPE_CHECKING:
@@ -25,10 +30,6 @@ if TYPE_CHECKING:
 
 logger: "logging.Logger" = get_logger("artifacts")
 
-if TYPE_CHECKING:
-    from prefect.client.orchestration import PrefectClient
-    from prefect.client.schemas.objects import Artifact as ArtifactResponse
-
 
 class Artifact(ArtifactRequest):
     """
@@ -43,10 +44,47 @@ class Artifact(ArtifactRequest):
         data: A JSON payload that allows for a result to be retrieved.
     """
 
-    @sync_compatible
-    async def create(
+    async def acreate(
+        self,
+        client: "PrefectClient | None" = None,
+    ) -> "ArtifactResponse":
+        """
+        An async method to create an artifact.
+
+        Arguments:
+            client: The PrefectClient
+
+        Returns:
+            - The created artifact.
+        """
+
+        local_client_context = asyncnullcontext(client) if client else get_client()
+        async with local_client_context as client:
+            task_run_id, flow_run_id = get_task_and_flow_run_ids()
+
+            try:
+                get_run_context()
+            except MissingContextError:
+                warnings.warn(
+                    "Artifact creation outside of a flow or task run is deprecated and will be removed in a later version.",
+                    FutureWarning,
+                )
+
+            return await client.create_artifact(
+                artifact=ArtifactRequest(
+                    type=self.type,
+                    key=self.key,
+                    description=self.description,
+                    task_run_id=self.task_run_id or task_run_id,
+                    flow_run_id=self.flow_run_id or flow_run_id,
+                    data=await self.aformat(),
+                )
+            )
+
+    @async_dispatch(acreate)
+    def create(
         self: Self,
-        client: Optional["PrefectClient"] = None,
+        client: "PrefectClient | None" = None,
     ) -> "ArtifactResponse":
         """
         A method to create an artifact.
@@ -57,9 +95,9 @@ class Artifact(ArtifactRequest):
         Returns:
             - The created artifact.
         """
-        from prefect.context import MissingContextError, get_run_context
 
-        client, _ = get_or_create_client(client)
+        # Create sync client since this is a sync method.
+        sync_client = get_client(sync_client=True)
         task_run_id, flow_run_id = get_task_and_flow_run_ids()
 
         try:
@@ -70,35 +108,67 @@ class Artifact(ArtifactRequest):
                 FutureWarning,
             )
 
-        return await client.create_artifact(
+        return sync_client.create_artifact(
             artifact=ArtifactRequest(
                 type=self.type,
                 key=self.key,
                 description=self.description,
                 task_run_id=self.task_run_id or task_run_id,
                 flow_run_id=self.flow_run_id or flow_run_id,
-                data=await self.format(),
+                data=cast(str, self.format(_sync=True)),  # pyright: ignore[reportCallIssue] _sync is valid because .format is wrapped in async_dispatch
             )
         )
 
     @classmethod
-    @sync_compatible
-    async def get(
-        cls, key: Optional[str] = None, client: Optional["PrefectClient"] = None
-    ) -> Optional["ArtifactResponse"]:
+    async def aget(
+        cls,
+        key: str | None = None,
+        client: "PrefectClient | None" = None,
+    ) -> "ArtifactResponse | None":
+        """
+        A async method to get an artifact.
+
+        Arguments:
+            key: The key of the artifact to get.
+            client: A client to use when calling the Prefect API.
+
+        Returns:
+            The artifact (if found).
+        """
+
+        local_client_context = asyncnullcontext(client) if client else get_client()
+        async with local_client_context as client:
+            filter_key_value = None if key is None else [key]
+            artifacts = await client.read_artifacts(
+                limit=1,
+                sort=ArtifactSort.UPDATED_DESC,
+                artifact_filter=ArtifactFilter(
+                    key=ArtifactFilterKey(any_=filter_key_value)
+                ),
+            )
+            return None if not artifacts else artifacts[0]
+
+    @classmethod
+    @async_dispatch(aget)
+    def get(
+        cls, key: str | None = None, client: "PrefectClient | None" = None
+    ) -> "ArtifactResponse | None":
         """
         A method to get an artifact.
 
         Arguments:
-            key (str, optional): The key of the artifact to get.
-            client (PrefectClient, optional): The PrefectClient
+            key: The key of the artifact to get.
+            client: A client to use when calling the Prefect API.
 
         Returns:
-            (ArtifactResponse, optional): The artifact (if found).
+            The artifact (if found).
         """
-        client, _ = get_or_create_client(client)
+
+        # Create sync client since this is a sync method.
+        sync_client = get_client(sync_client=True)
+
         filter_key_value = None if key is None else [key]
-        artifacts = await client.read_artifacts(
+        artifacts = sync_client.read_artifacts(
             limit=1,
             sort=ArtifactSort.UPDATED_DESC,
             artifact_filter=ArtifactFilter(
@@ -108,41 +178,75 @@ class Artifact(ArtifactRequest):
         return None if not artifacts else artifacts[0]
 
     @classmethod
-    @sync_compatible
-    async def get_or_create(
+    async def aget_or_create(
         cls,
-        key: Optional[str] = None,
-        description: Optional[str] = None,
-        data: Optional[Union[dict[str, Any], Any]] = None,
-        client: Optional["PrefectClient"] = None,
+        key: str | None = None,
+        description: str | None = None,
+        data: dict[str, Any] | Any | None = None,
+        client: "PrefectClient | None" = None,
+        **kwargs: Any,
+    ) -> tuple["ArtifactResponse", bool]:
+        """
+        A async method to get or create an artifact.
+
+        Arguments:
+            key: The key of the artifact to get or create.
+            description: The description of the artifact to create.
+            data: The data of the artifact to create.
+            client: The PrefectClient
+            **kwargs: Additional keyword arguments to use when creating the artifact.
+
+        Returns:
+            The artifact, either retrieved or created.
+        """
+        artifact = await cls.aget(key, client)
+        if artifact:
+            return artifact, False
+
+        new_artifact = cls(key=key, description=description, data=data, **kwargs)
+        created_artifact = await new_artifact.acreate(client)
+        return created_artifact, True
+
+    @classmethod
+    @async_dispatch(aget_or_create)
+    def get_or_create(
+        cls,
+        key: str | None = None,
+        description: str | None = None,
+        data: dict[str, Any] | Any | None = None,
+        client: "PrefectClient | None" = None,
         **kwargs: Any,
     ) -> tuple["ArtifactResponse", bool]:
         """
         A method to get or create an artifact.
 
         Arguments:
-            key (str, optional): The key of the artifact to get or create.
-            description (str, optional): The description of the artifact to create.
-            data (Union[Dict[str, Any], Any], optional): The data of the artifact to create.
-            client (PrefectClient, optional): The PrefectClient
+            key: The key of the artifact to get or create.
+            description: The description of the artifact to create.
+            data: The data of the artifact to create.
+            client: The PrefectClient
+            **kwargs: Additional keyword arguments to use when creating the artifact.
 
         Returns:
-            (ArtifactResponse): The artifact, either retrieved or created.
+            The artifact, either retrieved or created.
         """
-        artifact_coro = cls.get(key, client)
-        if TYPE_CHECKING:
-            assert asyncio.iscoroutine(artifact_coro)
-        artifact = await artifact_coro
+        artifact = cast(ArtifactResponse, cls.get(key, _sync=True))  # pyright: ignore[reportCallIssue] _sync is valid because .get is wrapped in async_dispatch
         if artifact:
             return artifact, False
 
         new_artifact = cls(key=key, description=description, data=data, **kwargs)
-        create_coro = new_artifact.create(client)
-        if TYPE_CHECKING:
-            assert asyncio.iscoroutine(create_coro)
-        return await create_coro, True
+        created_artifact = cast(
+            ArtifactResponse,
+            new_artifact.create(_sync=True),  # pyright: ignore[reportCallIssue] _sync is valid because .create is wrapped in async_dispatch
+        )
+        return created_artifact, True
+
+    # TODO: Remove this when we remove async_dispatch because it doesn't need to be async
+    async def aformat(self) -> str | float | int | dict[str, Any]:
+        return json.dumps(self.data)
 
-    async def format(self) -> Optional[Union[dict[str, Any], Any]]:
+    @async_dispatch(aformat)
+    def format(self) -> str | float | int | dict[str, Any]:
         return json.dumps(self.data)
 
 
@@ -151,19 +255,30 @@ class LinkArtifact(Artifact):
     link_text: Optional[str] = None
     type: Optional[str] = "markdown"
 
-    async def format(self) -> str:
+    def _format(self) -> str:
         return (
             f"[{self.link_text}]({self.link})"
             if self.link_text
             else f"[{self.link}]({self.link})"
         )
 
+    async def aformat(self) -> str:
+        return self._format()
+
+    @async_dispatch(aformat)
+    def format(self) -> str:
+        return self._format()
+
 
 class MarkdownArtifact(Artifact):
     markdown: str
     type: Optional[str] = "markdown"
 
-    async def format(self) -> str:
+    async def aformat(self) -> str:
+        return self.markdown
+
+    @async_dispatch(aformat)
+    def format(self) -> str:
         return self.markdown
 
 
@@ -173,8 +288,8 @@ class TableArtifact(Artifact):
 
     @classmethod
     def _sanitize(
-        cls, item: Union[dict[str, Any], list[Any], float]
-    ) -> Union[dict[str, Any], list[Any], int, float, None]:
+        cls, item: dict[str, Any] | list[Any] | float
+    ) -> dict[str, Any] | list[Any] | int | float | None:
         """
         Sanitize NaN values in a given item.
         The item can be a dict, list or float.
@@ -188,7 +303,11 @@ class TableArtifact(Artifact):
         else:
             return item
 
-    async def format(self) -> str:
+    async def aformat(self) -> str:
+        return json.dumps(self._sanitize(self.table))
+
+    @async_dispatch(aformat)
+    def format(self) -> str:
         return json.dumps(self._sanitize(self.table))
 
 
@@ -196,7 +315,7 @@ class ProgressArtifact(Artifact):
     progress: float
     type: Optional[str] = "progress"
 
-    async def format(self) -> float:
+    def _format(self) -> float:
         # Ensure progress is between 0 and 100
         min_progress = 0.0
         max_progress = 100.0
@@ -209,6 +328,13 @@ class ProgressArtifact(Artifact):
 
         return self.progress
 
+    async def aformat(self) -> float:
+        return self._format()
+
+    @async_dispatch(aformat)
+    def format(self) -> float:
+        return self._format()
+
 
 class ImageArtifact(Artifact):
     """
@@ -221,11 +347,14 @@ class ImageArtifact(Artifact):
     image_url: str
     type: Optional[str] = "image"
 
-    async def format(self) -> str:
+    async def aformat(self) -> str:
+        return self.image_url
+
+    @async_dispatch(aformat)
+    def format(self) -> str:
         """
         This method is used to format the artifact data so it can be properly sent
-        to the API when the .create() method is called. It is async because the
-        method is awaited in the parent class.
+        to the API when the .create() method is called.
 
         Returns:
             str: The image URL.
@@ -233,13 +362,46 @@ class ImageArtifact(Artifact):
         return self.image_url
 
 
-@sync_compatible
-async def create_link_artifact(
+async def acreate_link_artifact(
+    link: str,
+    link_text: str | None = None,
+    key: str | None = None,
+    description: str | None = None,
+    client: "PrefectClient | None" = None,
+) -> UUID:
+    """
+    Create a link artifact.
+
+    Arguments:
+        link: The link to create.
+        link_text: The link text.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+
+    Returns:
+        The table artifact ID.
+    """
+    new_artifact = LinkArtifact(
+        key=key,
+        description=description,
+        link=link,
+        link_text=link_text,
+    )
+    artifact = await new_artifact.acreate(client)
+
+    return artifact.id
+
+
+@async_dispatch(acreate_link_artifact)
+def create_link_artifact(
     link: str,
-    link_text: Optional[str] = None,
-    key: Optional[str] = None,
-    description: Optional[str] = None,
-    client: Optional["PrefectClient"] = None,
+    link_text: str | None = None,
+    key: str | None = None,
+    description: str | None = None,
+    client: "PrefectClient | None" = None,
 ) -> UUID:
     """
     Create a link artifact.
@@ -262,19 +424,44 @@ async def create_link_artifact(
         link=link,
         link_text=link_text,
     )
-    create_coro = new_artifact.create(client)
-    if TYPE_CHECKING:
-        assert asyncio.iscoroutine(create_coro)
-    artifact = await create_coro
+    artifact = cast(ArtifactResponse, new_artifact.create(_sync=True))  # pyright: ignore[reportCallIssue] _sync is valid because .create is wrapped in async_dispatch
+
+    return artifact.id
+
+
+async def acreate_markdown_artifact(
+    markdown: str,
+    key: str | None = None,
+    description: str | None = None,
+) -> UUID:
+    """
+    Create a markdown artifact.
+
+    Arguments:
+        markdown: The markdown to create.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+    Returns:
+        The table artifact ID.
+    """
+    new_artifact = MarkdownArtifact(
+        key=key,
+        description=description,
+        markdown=markdown,
+    )
+    artifact = await new_artifact.acreate()
 
     return artifact.id
 
 
-@sync_compatible
-async def create_markdown_artifact(
+@async_dispatch(acreate_markdown_artifact)
+def create_markdown_artifact(
     markdown: str,
-    key: Optional[str] = None,
-    description: Optional[str] = None,
+    key: str | None = None,
+    description: str | None = None,
 ) -> UUID:
     """
     Create a markdown artifact.
@@ -294,19 +481,45 @@ async def create_markdown_artifact(
         description=description,
         markdown=markdown,
     )
-    create_coro = new_artifact.create()
-    if TYPE_CHECKING:
-        assert asyncio.iscoroutine(create_coro)
-    artifact = await create_coro
+    artifact = cast(ArtifactResponse, new_artifact.create(_sync=True))  # pyright: ignore[reportCallIssue] _sync is valid because .create is wrapped in async_dispatch
 
     return artifact.id
 
 
-@sync_compatible
-async def create_table_artifact(
-    table: Union[dict[str, list[Any]], list[dict[str, Any]], list[list[Any]]],
-    key: Optional[str] = None,
-    description: Optional[str] = None,
+async def acreate_table_artifact(
+    table: dict[str, list[Any]] | list[dict[str, Any]] | list[list[Any]],
+    key: str | None = None,
+    description: str | None = None,
+) -> UUID:
+    """
+    Create a table artifact asynchronously.
+
+    Arguments:
+        table: The table to create.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+    Returns:
+        The table artifact ID.
+    """
+
+    new_artifact = TableArtifact(
+        key=key,
+        description=description,
+        table=table,
+    )
+    artifact = await new_artifact.acreate()
+
+    return artifact.id
+
+
+@async_dispatch(acreate_table_artifact)
+def create_table_artifact(
+    table: dict[str, list[Any]] | list[dict[str, Any]] | list[list[Any]],
+    key: str | None = None,
+    description: str | None = None,
 ) -> UUID:
     """
     Create a table artifact.
@@ -327,19 +540,45 @@ async def create_table_artifact(
         description=description,
         table=table,
     )
-    create_coro = new_artifact.create()
-    if TYPE_CHECKING:
-        assert asyncio.iscoroutine(create_coro)
-    artifact = await create_coro
+    artifact = cast(ArtifactResponse, new_artifact.create(_sync=True))  # pyright: ignore[reportCallIssue] _sync is valid because .create is wrapped in async_dispatch
 
     return artifact.id
 
 
-@sync_compatible
-async def create_progress_artifact(
+async def acreate_progress_artifact(
     progress: float,
-    key: Optional[str] = None,
-    description: Optional[str] = None,
+    key: str | None = None,
+    description: str | None = None,
+) -> UUID:
+    """
+    Create a progress artifact asynchronously.
+
+    Arguments:
+        progress: The percentage of progress represented by a float between 0 and 100.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+    Returns:
+        The progress artifact ID.
+    """
+
+    new_artifact = ProgressArtifact(
+        key=key,
+        description=description,
+        progress=progress,
+    )
+    artifact = await new_artifact.acreate()
+
+    return artifact.id
+
+
+@async_dispatch(acreate_progress_artifact)
+def create_progress_artifact(
+    progress: float,
+    key: str | None = None,
+    description: str | None = None,
 ) -> UUID:
     """
     Create a progress artifact.
@@ -360,20 +599,58 @@ async def create_progress_artifact(
         description=description,
         progress=progress,
     )
-    create_coro = new_artifact.create()
-    if TYPE_CHECKING:
-        assert asyncio.iscoroutine(create_coro)
-    artifact = await create_coro
+    artifact = cast(ArtifactResponse, new_artifact.create(_sync=True))  # pyright: ignore[reportCallIssue] _sync is valid because .create is wrapped in async_dispatch
 
     return artifact.id
 
 
-@sync_compatible
-async def update_progress_artifact(
+async def aupdate_progress_artifact(
+    artifact_id: UUID,
+    progress: float,
+    description: str | None = None,
+    client: "PrefectClient | None" = None,
+) -> UUID:
+    """
+    Update a progress artifact asynchronously.
+
+    Arguments:
+        artifact_id: The ID of the artifact to update.
+        progress: The percentage of progress represented by a float between 0 and 100.
+        description: A user-specified description of the artifact.
+
+    Returns:
+        The progress artifact ID.
+    """
+
+    local_client_context = nullcontext(client) if client else get_client()
+    async with local_client_context as client:
+        artifact = ProgressArtifact(
+            description=description,
+            progress=progress,
+        )
+        update = (
+            ArtifactUpdate(
+                description=artifact.description,
+                data=await artifact.aformat(),
+            )
+            if description
+            else ArtifactUpdate(data=await artifact.aformat())
+        )
+
+        await client.update_artifact(
+            artifact_id=artifact_id,
+            artifact=update,
+        )
+
+        return artifact_id
+
+
+@async_dispatch(aupdate_progress_artifact)
+def update_progress_artifact(
     artifact_id: UUID,
     progress: float,
-    description: Optional[str] = None,
-    client: Optional["PrefectClient"] = None,
+    description: str | None = None,
+    client: "PrefectClient | None" = None,
 ) -> UUID:
     """
     Update a progress artifact.
@@ -387,7 +664,7 @@ async def update_progress_artifact(
         The progress artifact ID.
     """
 
-    client, _ = get_or_create_client(client)
+    sync_client = get_client(sync_client=True)
 
     artifact = ProgressArtifact(
         description=description,
@@ -396,13 +673,13 @@ async def update_progress_artifact(
     update = (
         ArtifactUpdate(
             description=artifact.description,
-            data=await artifact.format(),
+            data=cast(float, artifact.format(_sync=True)),  # pyright: ignore[reportCallIssue] _sync is valid because .format is wrapped in async_dispatch
         )
         if description
-        else ArtifactUpdate(data=await artifact.format())
+        else ArtifactUpdate(data=cast(float, artifact.format(_sync=True)))  # pyright: ignore[reportCallIssue] _sync is valid because .format is wrapped in async_dispatch
     )
 
-    await client.update_artifact(
+    sync_client.update_artifact(
         artifact_id=artifact_id,
         artifact=update,
     )
@@ -410,11 +687,40 @@ async def update_progress_artifact(
     return artifact_id
 
 
-@sync_compatible
-async def create_image_artifact(
+async def acreate_image_artifact(
+    image_url: str,
+    key: str | None = None,
+    description: str | None = None,
+) -> UUID:
+    """
+    Create an image artifact asynchronously.
+
+    Arguments:
+        image_url: The URL of the image to display.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+    Returns:
+        The image artifact ID.
+    """
+
+    new_artifact = ImageArtifact(
+        key=key,
+        description=description,
+        image_url=image_url,
+    )
+    artifact = await new_artifact.acreate()
+
+    return artifact.id
+
+
+@async_dispatch(acreate_image_artifact)
+def create_image_artifact(
     image_url: str,
-    key: Optional[str] = None,
-    description: Optional[str] = None,
+    key: str | None = None,
+    description: str | None = None,
 ) -> UUID:
     """
     Create an image artifact.
@@ -435,9 +741,6 @@ async def create_image_artifact(
         description=description,
         image_url=image_url,
     )
-    create_coro = new_artifact.create()
-    if TYPE_CHECKING:
-        assert asyncio.iscoroutine(create_coro)
-    artifact = await create_coro
+    artifact = cast(ArtifactResponse, new_artifact.create(_sync=True))  # pyright: ignore[reportCallIssue] _sync is valid because .create is wrapped in async_dispatch
 
     return artifact.id