prefect-client 3.0.0rc1__py3-none-any.whl → 3.0.0rc3__py3-none-any.whl

prefect/flows.py

@@ -4,7 +4,6 @@ Module containing the base workflow class and decorator - for most use cases, us
 
 # This file requires type-checking with pyright because mypy does not yet support PEP612
 # See https://github.com/python/mypy/issues/8645
-
 import ast
 import datetime
 import importlib.util
@@ -15,6 +14,7 @@ import re
 import sys
 import tempfile
 import warnings
+from copy import copy
 from functools import partial, update_wrapper
 from pathlib import Path
 from tempfile import NamedTemporaryFile
@@ -102,7 +102,7 @@ from prefect.utilities.callables import (
 from prefect.utilities.collections import listrepr
 from prefect.utilities.filesystem import relative_path_to_current_platform
 from prefect.utilities.hashing import file_hash
-from prefect.utilities.importtools import import_object
+from prefect.utilities.importtools import import_object, safe_load_namespace
 
 from ._internal.pydantic.v2_schema import is_v2_type
 from ._internal.pydantic.v2_validated_func import V2ValidatedFunction
@@ -354,6 +354,28 @@ class Flow(Generic[P, R]):
 
         self._entrypoint = f"{module}:{fn.__name__}"
 
+    @property
+    def ismethod(self) -> bool:
+        return hasattr(self.fn, "__prefect_self__")
+
+    def __get__(self, instance, owner):
+        """
+        Implement the descriptor protocol so that the flow can be used as an instance method.
+        When an instance method is loaded, this method is called with the "self" instance as
+        an argument. We return a copy of the flow with that instance bound to the flow's function.
+        """
+
+        # if no instance is provided, it's being accessed on the class
+        if instance is None:
+            return self
+
+        # if the flow is being accessed on an instance, bind the instance to the __prefect_self__ attribute
+        # of the flow's function. This will allow it to be automatically added to the flow's parameters
+        else:
+            bound_flow = copy(self)
+            bound_flow.fn.__prefect_self__ = instance
+            return bound_flow
+
     def with_options(
         self,
         *,
@@ -555,6 +577,9 @@ class Flow(Generic[P, R]):
         """
         serialized_parameters = {}
         for key, value in parameters.items():
+            # do not serialize the bound self object
+            if self.ismethod and value is self.fn.__prefect_self__:
+                continue
             try:
                 serialized_parameters[key] = jsonable_encoder(value)
             except (TypeError, ValueError):
@@ -1241,19 +1266,14 @@ class Flow(Generic[P, R]):
             # we can add support for exploring subflows for tasks in the future.
             return track_viz_task(self.isasync, self.name, parameters)
 
-        from prefect.flow_engine import run_flow, run_flow_sync
+        from prefect.flow_engine import run_flow
 
-        run_kwargs = dict(
+        return run_flow(
             flow=self,
             parameters=parameters,
             wait_for=wait_for,
             return_type=return_type,
         )
-        if self.isasync:
-            # this returns an awaitable coroutine
-            return run_flow(**run_kwargs)
-        else:
-            return run_flow_sync(**run_kwargs)
 
     @sync_compatible
     async def visualize(self, *args, **kwargs):
@@ -1329,8 +1349,8 @@ def flow(
     retries: Optional[int] = None,
     retry_delay_seconds: Optional[Union[int, float]] = None,
     task_runner: Optional[TaskRunner] = None,
-    description: str = None,
-    timeout_seconds: Union[int, float] = None,
+    description: Optional[str] = None,
+    timeout_seconds: Union[int, float, None] = None,
     validate_parameters: bool = True,
     persist_result: Optional[bool] = None,
     result_storage: Optional[ResultStorage] = None,
@@ -1358,11 +1378,11 @@ def flow(
     name: Optional[str] = None,
     version: Optional[str] = None,
     flow_run_name: Optional[Union[Callable[[], str], str]] = None,
-    retries: int = None,
-    retry_delay_seconds: Union[int, float] = None,
+    retries: Optional[int] = None,
+    retry_delay_seconds: Union[int, float, None] = None,
     task_runner: Optional[TaskRunner] = None,
-    description: str = None,
-    timeout_seconds: Union[int, float] = None,
+    description: Optional[str] = None,
+    timeout_seconds: Union[int, float, None] = None,
     validate_parameters: bool = True,
     persist_result: Optional[bool] = None,
     result_storage: Optional[ResultStorage] = None,
@@ -1485,6 +1505,9 @@ def flow(
         >>>     pass
     """
     if __fn:
+        if isinstance(__fn, (classmethod, staticmethod)):
+            method_decorator = type(__fn).__name__
+            raise TypeError(f"@{method_decorator} should be applied on top of @flow")
         return cast(
             Flow[P, R],
             Flow(
@@ -1560,7 +1583,9 @@ flow.from_source = Flow.from_source
 
 
 def select_flow(
-    flows: Iterable[Flow], flow_name: str = None, from_message: str = None
+    flows: Iterable[Flow],
+    flow_name: Optional[str] = None,
+    from_message: Optional[str] = None,
 ) -> Flow:
     """
     Select the only flow in an iterable or a flow specified by name.
@@ -1574,33 +1599,33 @@ def select_flow(
         UnspecifiedFlowError: If multiple flows exist but no flow name was provided
     """
     # Convert to flows by name
-    flows = {f.name: f for f in flows}
+    flows_dict = {f.name: f for f in flows}
 
     # Add a leading space if given, otherwise use an empty string
     from_message = (" " + from_message) if from_message else ""
-    if not flows:
+    if not Optional:
         raise MissingFlowError(f"No flows found{from_message}.")
 
-    elif flow_name and flow_name not in flows:
+    elif flow_name and flow_name not in flows_dict:
         raise MissingFlowError(
             f"Flow {flow_name!r} not found{from_message}. "
-            f"Found the following flows: {listrepr(flows.keys())}. "
+            f"Found the following flows: {listrepr(flows_dict.keys())}. "
             "Check to make sure that your flow function is decorated with `@flow`."
         )
 
-    elif not flow_name and len(flows) > 1:
+    elif not flow_name and len(flows_dict) > 1:
         raise UnspecifiedFlowError(
             (
-                f"Found {len(flows)} flows{from_message}:"
-                f" {listrepr(sorted(flows.keys()))}. Specify a flow name to select a"
+                f"Found {len(flows_dict)} flows{from_message}:"
+                f" {listrepr(sorted(flows_dict.keys()))}. Specify a flow name to select a"
                 " flow."
             ),
         )
 
     if flow_name:
-        return flows[flow_name]
+        return flows_dict[flow_name]
     else:
-        return list(flows.values())[0]
+        return list(flows_dict.values())[0]
 
 
 def load_flows_from_script(path: str) -> List[Flow]:
@@ -1617,7 +1642,7 @@ def load_flows_from_script(path: str) -> List[Flow]:
     return registry_from_script(path).get_instances(Flow)
 
 
-def load_flow_from_script(path: str, flow_name: str = None) -> Flow:
+def load_flow_from_script(path: str, flow_name: Optional[str] = None) -> Flow:
     """
     Extract a flow object from a script by running all of the code in the file.
 
@@ -1661,7 +1686,7 @@ def load_flow_from_entrypoint(
         FlowScriptError: If an exception is encountered while running the script
         MissingFlowError: If the flow function specified in the entrypoint does not exist
     """
-    with PrefectObjectRegistry(
+    with PrefectObjectRegistry(  # type: ignore
         block_code_execution=True,
         capture_failures=True,
     ):
@@ -1686,7 +1711,7 @@ def load_flow_from_entrypoint(
         return flow
 
 
-def load_flow_from_text(script_contents: AnyStr, flow_name: str):
+def load_flow_from_text(script_contents: AnyStr, flow_name: str) -> Flow:
     """
     Load a flow from a text script.
 
@@ -1717,7 +1742,7 @@ async def serve(
     print_starting_message: bool = True,
     limit: Optional[int] = None,
     **kwargs,
-):
+) -> NoReturn:
     """
     Serve the provided list of deployments.
 
@@ -1807,7 +1832,7 @@ async def load_flow_from_flow_run(
     flow_run: "FlowRun",
     ignore_storage: bool = False,
     storage_base_path: Optional[str] = None,
-) -> "Flow":
+) -> Flow:
     """
     Load a flow from the location/script provided in a deployment's storage document.
 
@@ -1861,7 +1886,9 @@ async def load_flow_from_flow_run(
         await storage_block.get_directory(from_path=from_path, local_path=".")
 
     if deployment.pull_steps:
-        run_logger.debug(f"Running {len(deployment.pull_steps)} deployment pull steps")
+        run_logger.debug(
+            f"Running {len(deployment.pull_steps)} deployment pull step(s)"
+        )
         output = await run_steps(deployment.pull_steps)
         if output.get("directory"):
             run_logger.debug(f"Changing working directory to {output['directory']!r}")
@@ -1913,7 +1940,14 @@ def load_flow_argument_from_entrypoint(
         (
             node
             for node in ast.walk(parsed_code)
-            if isinstance(node, ast.FunctionDef) and node.name == func_name
+            if isinstance(
+                node,
+                (
+                    ast.FunctionDef,
+                    ast.AsyncFunctionDef,
+                ),
+            )
+            and node.name == func_name
         ),
         None,
     )
@@ -1926,11 +1960,33 @@ def load_flow_argument_from_entrypoint(
         ):
             for keyword in decorator.keywords:
                 if keyword.arg == arg:
-                    return (
-                        keyword.value.value
-                    )  # Return the string value of the argument
+                    if isinstance(keyword.value, ast.Constant):
+                        return (
+                            keyword.value.value
+                        )  # Return the string value of the argument
+
+                    # if the arg value is not a raw str (i.e. a variable or expression),
+                    # then attempt to evaluate it
+                    namespace = safe_load_namespace(source_code)
+                    literal_arg_value = ast.get_source_segment(
+                        source_code, keyword.value
+                    )
+                    try:
+                        evaluated_value = eval(literal_arg_value, namespace)  # type: ignore
+                    except Exception as e:
+                        logger.info(
+                            "Failed to parse @flow argument: `%s=%s` due to the following error. Ignoring and falling back to default behavior.",
+                            arg,
+                            literal_arg_value,
+                            exc_info=e,
+                        )
+                        # ignore the decorator arg and fallback to default behavior
+                        break
+                    return str(evaluated_value)
 
     if arg == "name":
         return func_name.replace(
             "_", "-"
         )  # If no matching decorator or keyword argument is found
+
+    return None

prefect/futures.py

@@ -1,7 +1,6 @@
 import abc
 import concurrent.futures
 import inspect
-import time
 import uuid
 from functools import partial
 from typing import Any, Generic, Optional, Set, Union, cast
@@ -11,13 +10,17 @@ from typing_extensions import TypeVar
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.objects import TaskRun
 from prefect.exceptions import ObjectNotFound
+from prefect.logging.loggers import get_logger
 from prefect.states import Pending, State
+from prefect.task_runs import TaskRunWaiter
 from prefect.utilities.annotations import quote
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import StopVisiting, visit_collection
 
 F = TypeVar("F")
 
+logger = get_logger(__name__)
+
 
 class PrefectFuture(abc.ABC):
     """
@@ -146,68 +149,66 @@ class PrefectDistributedFuture(PrefectFuture):
     Represents the result of a computation happening anywhere.
 
     This class is typically used to interact with the result of a task run
-    scheduled to run in a Prefect task server but can be used to interact with
+    scheduled to run in a Prefect task worker but can be used to interact with
     any task run scheduled in Prefect's API.
     """
 
-    def __init__(self, task_run_id: uuid.UUID):
-        self._task_run = None
-        self._client = None
-        super().__init__(task_run_id=task_run_id)
-
-    @property
-    def client(self):
-        if self._client is None:
-            self._client = get_client(sync_client=True)
-        return self._client
+    def wait(self, timeout: Optional[float] = None) -> None:
+        return run_coro_as_sync(self.wait_async(timeout=timeout))
 
-    @property
-    def task_run(self):
-        if self._task_run is None:
-            self._task_run = self.client.read_task_run(task_run_id=self.task_run_id)
-        return self._task_run
-
-    @task_run.setter
-    def task_run(self, task_run):
-        self._task_run = task_run
-
-    def wait(
-        self, timeout: Optional[float] = None, polling_interval: Optional[float] = 0.2
-    ) -> None:
-        start_time = time.time()
-        # TODO: Websocket implementation?
-        while True:
-            self.task_run = cast(
-                TaskRun, self.client.read_task_run(task_run_id=self.task_run_id)
+    async def wait_async(self, timeout: Optional[float] = None):
+        if self._final_state:
+            logger.debug(
+                "Final state already set for %s. Returning...", self.task_run_id
             )
-            if self.task_run.state and self.task_run.state.is_final():
-                self._final_state = self.task_run.state
-                return
-            if timeout is not None and (time.time() - start_time) > timeout:
+            return
+
+        # Read task run to see if it is still running
+        async with get_client() as client:
+            task_run = await client.read_task_run(task_run_id=self._task_run_id)
+            if task_run.state.is_final():
+                logger.debug(
+                    "Task run %s already finished. Returning...",
+                    self.task_run_id,
+                )
+                self._final_state = task_run.state
                 return
-            time.sleep(polling_interval)
+
+            # If still running, wait for a completed event from the server
+            logger.debug(
+                "Waiting for completed event for task run %s...",
+                self.task_run_id,
+            )
+            await TaskRunWaiter.wait_for_task_run(self._task_run_id, timeout=timeout)
+            task_run = await client.read_task_run(task_run_id=self._task_run_id)
+            if task_run.state.is_final():
+                self._final_state = task_run.state
+            return
 
     def result(
         self,
         timeout: Optional[float] = None,
         raise_on_failure: bool = True,
-        polling_interval: Optional[float] = 0.2,
     ) -> Any:
+        return run_coro_as_sync(
+            self.result_async(timeout=timeout, raise_on_failure=raise_on_failure)
+        )
+
+    async def result_async(
+        self,
+        timeout: Optional[float] = None,
+        raise_on_failure: bool = True,
+    ):
         if not self._final_state:
-            self.wait(timeout=timeout)
+            await self.wait_async(timeout=timeout)
             if not self._final_state:
                 raise TimeoutError(
                     f"Task run {self.task_run_id} did not complete within {timeout} seconds"
                 )
 
-        _result = self._final_state.result(
+        return await self._final_state.result(
             raise_on_failure=raise_on_failure, fetch=True
         )
-        # state.result is a `sync_compatible` function that may or may not return an awaitable
-        # depending on whether the parent frame is sync or not
-        if inspect.isawaitable(_result):
-            _result = run_coro_as_sync(_result)
-        return _result
 
     def __eq__(self, other):
         if not isinstance(other, PrefectDistributedFuture):

prefect/infrastructure/provisioners/container_instance.py

@@ -10,6 +10,7 @@ Classes:
     ContainerInstancePushProvisioner: A class for provisioning infrastructure using Azure Container Instances.
 
 """
+
 import json
 import random
 import shlex

prefect/infrastructure/provisioners/ecs.py

@@ -367,7 +367,7 @@ class AuthenticationResource:
         work_pool_name: str,
         user_name: str = "prefect-ecs-user",
         policy_name: str = "prefect-ecs-policy",
-        credentials_block_name: str = None,
+        credentials_block_name: Optional[str] = None,
     ):
         self._user_name = user_name
         self._credentials_block_name = (
@@ -1130,7 +1130,7 @@ class ElasticContainerServicePushProvisioner:
         work_pool_name: str,
         user_name: str = "prefect-ecs-user",
         policy_name: str = "prefect-ecs-policy",
-        credentials_block_name: str = None,
+        credentials_block_name: Optional[str] = None,
         cluster_name: str = "prefect-ecs-cluster",
         vpc_name: str = "prefect-ecs-vpc",
         ecs_security_group_name: str = "prefect-ecs-security-group",

prefect/input/__init__.py

@@ -12,6 +12,8 @@ from .run_input import (
     RunInputMetadata,
     keyset_from_base_key,
     keyset_from_paused_state,
+    receive_input,
+    send_input,
 )
 
 __all__ = [
@@ -26,4 +28,6 @@ __all__ = [
     "keyset_from_base_key",
     "keyset_from_paused_state",
     "read_flow_run_input",
+    "receive_input",
+    "send_input",
 ]

prefect/input/run_input.py

@@ -18,7 +18,8 @@ Sender flow:
 ```python
 import random
 from uuid import UUID
-from prefect import flow, get_run_logger
+from prefect import flow
+from prefect.logging import get_run_logger
 from prefect.input import RunInput
 
 class NumberData(RunInput):
@@ -43,7 +44,8 @@ Receiver flow:
 ```python
 import random
 from uuid import UUID
-from prefect import flow, get_run_logger
+from prefect import flow
+from prefect.logging import get_run_logger
 from prefect.input import RunInput
 
 class NumberData(RunInput):

prefect/logging/formatters.py

@@ -78,8 +78,8 @@ class PrefectFormatter(logging.Formatter):
         validate=True,
         *,
         defaults=None,
-        task_run_fmt: str = None,
-        flow_run_fmt: str = None,
+        task_run_fmt: Optional[str] = None,
+        flow_run_fmt: Optional[str] = None,
     ) -> None:
         """
         Implementation of the standard Python formatter with support for multiple

prefect/logging/handlers.py

@@ -108,8 +108,8 @@ class APILogHandler(logging.Handler):
                 )
 
             # Not ideal, but this method is called by the stdlib and cannot return a
-            # coroutine so we just schedule the drain in a new thread and continue
-            from_sync.call_soon_in_new_thread(create_call(APILogWorker.drain_all))
+            # coroutine so we just schedule the drain in the global loop thread and continue
+            from_sync.call_soon_in_loop_thread(create_call(APILogWorker.drain_all))
             return None
         else:
             # We set a timeout of 5s because we don't want to block forever if the worker

prefect/logging/loggers.py

@@ -69,7 +69,7 @@ class PrefectLogAdapter(logging.LoggerAdapter):
 
 
 @lru_cache()
-def get_logger(name: str = None) -> logging.Logger:
+def get_logger(name: Optional[str] = None) -> logging.Logger:
     """
     Get a `prefect` logger. These loggers are intended for internal use within the
     `prefect` package.

prefect/plugins.py

@@ -7,6 +7,7 @@ Currently supported entrypoints:
     - prefect.collections: Identifies this package as a Prefect collection that
         should be imported when Prefect is imported.
 """
+
 import sys
 from types import ModuleType
 from typing import Any, Dict, Union