prefect 3.1.11.dev2__tar.gz → 3.1.11.dev3__tar.gz

{prefect-3.1.11.dev2/src/prefect.egg-info → prefect-3.1.11.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect
-Version: 3.1.11.dev2
+Version: 3.1.11.dev3
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.

{prefect-3.1.11.dev2 → prefect-3.1.11.dev3}/src/prefect/_experimental/lineage.py

@@ -86,18 +86,18 @@ async def emit_lineage_event(
     else:
         upstream_resources.extend(related_resources)
 
+    # We want to consider all resources upstream and downstream of the event as
+    # lineage-related, including flows, flow runs, etc., so we add the label to
+    # all resources involved in the event.
+    for res in upstream_resources + downstream_resources:
+        if "prefect.resource.lineage-group" not in res:
+            res["prefect.resource.lineage-group"] = "global"
+
     # Emit an event for each downstream resource. This is necessary because
     # our event schema allows one primary resource and many related resources,
     # and for the purposes of lineage, related resources can only represent
     # upstream resources.
     for resource in downstream_resources:
-        # Downstream lineage resources need to have the
-        # prefect.resource.lineage-group label. All upstram resources from a
-        # downstream resource with this label will be considered lineage-related
-        # resources.
-        if "prefect.resource.lineage-group" not in resource:
-            resource["prefect.resource.lineage-group"] = "global"
-
         emit_kwargs: Dict[str, Any] = {
             "event": event_name,
             "resource": resource,
@@ -170,7 +170,6 @@ async def emit_result_write_event(
             {
                 "prefect.resource.id": result_resource_uri,
                 "prefect.resource.role": "result",
-                "prefect.resource.lineage-group": "global",
             }
         ]
         await emit_lineage_event(

{prefect-3.1.11.dev2 → prefect-3.1.11.dev3}/src/prefect/_internal/compatibility/deprecated.py

@@ -56,7 +56,7 @@ def generate_deprecation_message(
     end_date: Optional[str] = None,
     help: str = "",
     when: str = "",
-):
+) -> str:
     if not start_date and not end_date:
         raise ValueError(
             "A start date is required if an end date is not provided. Suggested start"

{prefect-3.1.11.dev2 → prefect-3.1.11.dev3}/src/prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-28T18:21:09-0800",
+ "date": "2024-12-30T22:24:19-0600",
  "dirty": false,
  "error": null,
- "full-revisionid": "2133d67536e63dc09728e5ba50b4fa227977f651",
- "version": "3.1.11.dev2"
+ "full-revisionid": "7e6f88322698f87cc2b495c370110bf46143cfa1",
+ "version": "3.1.11.dev3"
 }
 '''  # END VERSION_JSON
 

{prefect-3.1.11.dev2 → prefect-3.1.11.dev3}/src/prefect/blocks/core.py

@@ -38,7 +38,7 @@ from pydantic import (
     model_serializer,
 )
 from pydantic.json_schema import GenerateJsonSchema
-from typing_extensions import Literal, ParamSpec, Self, get_args
+from typing_extensions import Literal, ParamSpec, Self, TypeGuard, get_args
 
 import prefect.exceptions
 from prefect._internal.compatibility.async_dispatch import async_dispatch
@@ -1163,7 +1163,7 @@ class Block(BaseModel, ABC):
             ) from e
 
     @staticmethod
-    def is_block_class(block) -> bool:
+    def is_block_class(block: Any) -> TypeGuard[Type["Block"]]:
         return _is_subclass(block, Block)
 
     @staticmethod

{prefect-3.1.11.dev2 → prefect-3.1.11.dev3}/src/prefect/cli/_prompts.py

@@ -2,14 +2,28 @@
 Utilities for prompting the user for input
 """
 
+from __future__ import annotations
+
+import ast
+import asyncio
+import math
 import os
 import shutil
 from datetime import timedelta
 from getpass import GetPassWarning
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Coroutine,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
 
+import anyio
 import readchar
-from rich.console import Console, Group
+from rich.console import Console, Group, RenderableType
 from rich.live import Live
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.prompt import Confirm, InvalidResponse, Prompt, PromptBase
@@ -26,21 +40,25 @@ from prefect.client.schemas.schedules import (
     CronSchedule,
     IntervalSchedule,
     RRuleSchedule,
+    is_valid_timezone,
 )
 from prefect.client.utilities import client_injector
-from prefect.deployments.base import (
-    _get_git_remote_origin_url,
-    _search_for_flow_functions,
-)
 from prefect.exceptions import ObjectAlreadyExists, ObjectNotFound
 from prefect.flows import load_flow_from_entrypoint
+from prefect.logging.loggers import get_logger
+from prefect.settings import PREFECT_DEBUG_MODE
 from prefect.utilities import urls
+from prefect.utilities._git import get_git_remote_origin_url
+from prefect.utilities.asyncutils import LazySemaphore
+from prefect.utilities.filesystem import get_open_file_limit
 from prefect.utilities.processutils import get_sys_executable, run_process
 from prefect.utilities.slugify import slugify
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
 
+T = TypeVar("T", bound=RenderableType)
+
 STORAGE_PROVIDER_TO_CREDS_BLOCK = {
     "s3": "aws-credentials",
     "gcs": "gcp-credentials",
@@ -53,26 +71,158 @@ REQUIRED_FIELDS_FOR_CREDS_BLOCK = {
     "azure-blob-storage-credentials": ["account_url", "connection_string"],
 }
 
+# Only allow half of the open file limit to be open at once to allow for other
+# actors to open files.
+OPEN_FILE_SEMAPHORE = LazySemaphore(lambda: math.floor(get_open_file_limit() * 0.5))
+
+
+async def find_flow_functions_in_file(path: anyio.Path) -> list[dict[str, str]]:
+    decorator_name = "flow"
+    decorator_module = "prefect"
+    decorated_functions: list[dict[str, str]] = []
+    async with OPEN_FILE_SEMAPHORE:
+        try:
+            async with await anyio.open_file(path) as f:
+                try:
+                    tree = ast.parse(await f.read())
+                except SyntaxError:
+                    if PREFECT_DEBUG_MODE:
+                        get_logger().debug(
+                            f"Could not parse {path} as a Python file. Skipping."
+                        )
+                    return decorated_functions
+        except Exception as exc:
+            if PREFECT_DEBUG_MODE:
+                get_logger().debug(f"Could not open {path}: {exc}. Skipping.")
+            return decorated_functions
+
+    for node in ast.walk(tree):
+        if isinstance(
+            node,
+            (
+                ast.FunctionDef,
+                ast.AsyncFunctionDef,
+            ),
+        ):
+            for decorator in node.decorator_list:
+                # handles @flow
+                is_name_match = (
+                    isinstance(decorator, ast.Name) and decorator.id == decorator_name
+                )
+                # handles @flow()
+                is_func_name_match = (
+                    isinstance(decorator, ast.Call)
+                    and isinstance(decorator.func, ast.Name)
+                    and decorator.func.id == decorator_name
+                )
+                # handles @prefect.flow
+                is_module_attribute_match = (
+                    isinstance(decorator, ast.Attribute)
+                    and isinstance(decorator.value, ast.Name)
+                    and decorator.value.id == decorator_module
+                    and decorator.attr == decorator_name
+                )
+                # handles @prefect.flow()
+                is_module_attribute_func_match = (
+                    isinstance(decorator, ast.Call)
+                    and isinstance(decorator.func, ast.Attribute)
+                    and decorator.func.attr == decorator_name
+                    and isinstance(decorator.func.value, ast.Name)
+                    and decorator.func.value.id == decorator_module
+                )
+                if is_name_match or is_module_attribute_match:
+                    decorated_functions.append(
+                        {
+                            "flow_name": node.name,
+                            "function_name": node.name,
+                            "filepath": str(path),
+                        }
+                    )
+                if is_func_name_match or is_module_attribute_func_match:
+                    name_kwarg_node = None
+                    if TYPE_CHECKING:
+                        assert isinstance(decorator, ast.Call)
+                    for kw in decorator.keywords:
+                        if kw.arg == "name":
+                            name_kwarg_node = kw
+                            break
+                    if name_kwarg_node is not None and isinstance(
+                        name_kwarg_node.value, ast.Constant
+                    ):
+                        flow_name = name_kwarg_node.value.value
+                    else:
+                        flow_name = node.name
+                    decorated_functions.append(
+                        {
+                            "flow_name": flow_name,
+                            "function_name": node.name,
+                            "filepath": str(path),
+                        }
+                    )
+    return decorated_functions
+
+
+async def search_for_flow_functions(directory: str = ".") -> list[dict[str, str]]:
+    """
+    Search for flow functions in the provided directory. If no directory is provided,
+    the current working directory is used.
 
-def prompt(message, **kwargs):
+    Returns:
+        List[Dict]: the flow name, function name, and filepath of all flow functions found
+    """
+    path = anyio.Path(directory)
+    coros: list[Coroutine[list[dict[str, str]], Any, Any]] = []
+    async for file in path.rglob("*.py"):
+        coros.append(find_flow_functions_in_file(file))
+
+    return [fn for file_fns in await asyncio.gather(*coros) for fn in file_fns]
+
+
+def prompt(message: str, **kwargs: Any) -> str:
     """Utility to prompt the user for input with consistent styling"""
     return Prompt.ask(f"[bold][green]?[/] {message}[/]", **kwargs)
 
 
-def confirm(message, **kwargs):
+def confirm(message: str, **kwargs: Any) -> bool:
     """Utility to prompt the user for confirmation with consistent styling"""
     return Confirm.ask(f"[bold][green]?[/] {message}[/]", **kwargs)
 
 
+@overload
 def prompt_select_from_table(
-    console,
+    console: Console,
     prompt: str,
-    columns: List[Dict],
-    data: List[Dict],
-    table_kwargs: Optional[Dict] = None,
-    opt_out_message: Optional[str] = None,
+    columns: list[dict[str, str]],
+    data: list[dict[str, T]],
+    table_kwargs: dict[str, Any] | None = None,
+    opt_out_message: None = None,
     opt_out_response: Any = None,
-) -> Dict:
+) -> dict[str, T]:
+    ...
+
+
+@overload
+def prompt_select_from_table(
+    console: Console,
+    prompt: str,
+    columns: list[dict[str, str]],
+    data: list[dict[str, T]],
+    table_kwargs: dict[str, Any] | None = None,
+    opt_out_message: str = "",
+    opt_out_response: Any = None,
+) -> dict[str, T] | None:
+    ...
+
+
+def prompt_select_from_table(
+    console: Console,
+    prompt: str,
+    columns: list[dict[str, str]],
+    data: list[dict[str, T]],
+    table_kwargs: dict[str, Any] | None = None,
+    opt_out_message: str | None = None,
+    opt_out_response: Any = None,
+) -> dict[str, T] | None:
     """
     Given a list of columns and some data, display options to user in a table
     and prompt them to select one.
@@ -163,7 +313,7 @@ def prompt_select_from_table(
 
 # Interval schedule prompting utilities
 class IntervalValuePrompt(PromptBase[timedelta]):
-    response_type = timedelta
+    response_type: type[timedelta] = timedelta
     validate_error_message = (
         "[prompt.invalid]Please enter a valid interval denoted in seconds"
     )
@@ -178,7 +328,7 @@ class IntervalValuePrompt(PromptBase[timedelta]):
             raise InvalidResponse(self.validate_error_message)
 
 
-def prompt_interval_schedule(console):
+def prompt_interval_schedule(console: Console) -> IntervalSchedule:
     """
     Prompt the user for an interval in seconds.
     """
@@ -201,7 +351,7 @@ def prompt_interval_schedule(console):
 
 
 class CronStringPrompt(PromptBase[str]):
-    response_type = str
+    response_type: type[str] = str
     validate_error_message = "[prompt.invalid]Please enter a valid cron string"
 
     def process_response(self, value: str) -> str:
@@ -213,7 +363,7 @@ class CronStringPrompt(PromptBase[str]):
 
 
 class CronTimezonePrompt(PromptBase[str]):
-    response_type = str
+    response_type: type[str] = str
     validate_error_message = "[prompt.invalid]Please enter a valid timezone."
 
     def process_response(self, value: str) -> str:
@@ -224,7 +374,7 @@ class CronTimezonePrompt(PromptBase[str]):
             raise InvalidResponse(self.validate_error_message)
 
 
-def prompt_cron_schedule(console):
+def prompt_cron_schedule(console: Console) -> CronSchedule:
     """
     Prompt the user for a cron string and timezone.
     """
@@ -243,7 +393,7 @@ def prompt_cron_schedule(console):
 
 
 class RRuleStringPrompt(PromptBase[str]):
-    response_type = str
+    response_type: type[str] = str
     validate_error_message = "[prompt.invalid]Please enter a valid RRule string"
 
     def process_response(self, value: str) -> str:
@@ -255,18 +405,18 @@ class RRuleStringPrompt(PromptBase[str]):
 
 
 class RRuleTimezonePrompt(PromptBase[str]):
-    response_type = str
+    response_type: type[str] = str
     validate_error_message = "[prompt.invalid]Please enter a valid timezone."
 
     def process_response(self, value: str) -> str:
         try:
-            RRuleSchedule.valid_timezone(value)
+            is_valid_timezone(value)
             return value
         except ValueError:
             raise InvalidResponse(self.validate_error_message)
 
 
-def prompt_rrule_schedule(console):
+def prompt_rrule_schedule(console: Console) -> RRuleSchedule:
     """
     Prompts the user to enter an RRule string and timezone.
     """
@@ -284,7 +434,7 @@ def prompt_rrule_schedule(console):
 # Schedule type prompting utilities
 
 
-def prompt_schedule_type(console):
+def prompt_schedule_type(console: Console) -> str:
     """
     Prompts the user to select a schedule type from a list of options.
     """
@@ -322,11 +472,11 @@ def prompt_schedule_type(console):
     return selection["type"]
 
 
-def prompt_schedules(console) -> List[DeploymentScheduleCreate]:
+def prompt_schedules(console: Console) -> list[DeploymentScheduleCreate]:
     """
     Prompt the user to configure schedules for a deployment.
     """
-    schedules = []
+    schedules: list[DeploymentScheduleCreate] = []
 
     if confirm(
         "Would you like to configure schedules for this deployment?", default=True
@@ -346,12 +496,10 @@ def prompt_schedules(console) -> List[DeploymentScheduleCreate]:
             is_schedule_active = confirm(
                 "Would you like to activate this schedule?", default=True
             )
-            minimal_schedule_kwargs = {
-                "schedule": schedule,
-                "active": is_schedule_active,
-            }
 
-            schedules.append(DeploymentScheduleCreate(**minimal_schedule_kwargs))
+            schedules.append(
+                DeploymentScheduleCreate(schedule=schedule, active=is_schedule_active)
+            )
 
             add_schedule = confirm(
                 "Would you like to add another schedule?", default=False
@@ -391,8 +539,8 @@ async def prompt_select_work_pool(
 
 async def prompt_build_custom_docker_image(
     console: Console,
-    deployment_config: dict,
-):
+    deployment_config: dict[str, Any],
+) -> dict[str, Any] | None:
     if not confirm(
         "Would you like to build a custom Docker image for this deployment?",
         console=console,
@@ -451,9 +599,9 @@ async def prompt_build_custom_docker_image(
 
 async def prompt_push_custom_docker_image(
     console: Console,
-    deployment_config: dict,
-    build_docker_image_step: dict,
-):
+    deployment_config: dict[str, Any],
+    build_docker_image_step: dict[str, Any],
+) -> tuple[dict[str, Any] | None, dict[str, Any]]:
     if not confirm(
         "Would you like to push this image to a remote registry?",
         console=console,
@@ -537,10 +685,12 @@ async def prompt_push_custom_docker_image(
                     password=docker_credentials["password"],
                     registry_url=docker_credentials["registry_url"],
                 )
-                await new_creds_block.save(
+                coro = new_creds_block.save(
                     name=docker_registry_creds_name, overwrite=True
                 )
-
+                if TYPE_CHECKING:
+                    assert asyncio.iscoroutine(coro)
+                await coro
     return {
         "prefect_docker.deployments.steps.push_docker_image": push_step
     }, build_docker_image_step
@@ -589,7 +739,7 @@ async def prompt_create_work_pool(
 
 
 class EntrypointPrompt(PromptBase[str]):
-    response_type = str
+    response_type: type[str] = str
     validate_error_message = "[prompt.invalid]Please enter a valid flow entrypoint."
 
     def process_response(self, value: str) -> str:
@@ -624,7 +774,7 @@ async def prompt_entrypoint(console: Console) -> str:
             description="Scanning for flows...",
             total=1,
         )
-        discovered_flows = await _search_for_flow_functions()
+        discovered_flows = await search_for_flow_functions()
         progress.update(task_id, completed=1)
     if not discovered_flows:
         return EntrypointPrompt.ask(
@@ -659,7 +809,7 @@ async def prompt_entrypoint(console: Console) -> str:
 async def prompt_select_remote_flow_storage(
     client: "PrefectClient", console: Console
 ) -> Optional[str]:
-    valid_slugs_for_context = set()
+    valid_slugs_for_context: set[str] = set()
 
     for (
         storage_provider,
@@ -673,7 +823,7 @@ async def prompt_select_remote_flow_storage(
         except ObjectNotFound:
             pass
 
-    if _get_git_remote_origin_url():
+    if get_git_remote_origin_url():
         valid_slugs_for_context.add("git")
 
     flow_storage_options = [
@@ -700,9 +850,7 @@ async def prompt_select_remote_flow_storage(
     ]
 
     valid_storage_options_for_context = [
-        row
-        for row in flow_storage_options
-        if row is not None and row["slug"] in valid_slugs_for_context
+        row for row in flow_storage_options if row["slug"] in valid_slugs_for_context
     ]
 
     selected_flow_storage_row = prompt_select_from_table(
@@ -752,7 +900,11 @@ async def prompt_select_blob_storage_credentials(
                     "key": "name",
                 }
             ],
-            data=[{"name": block.name} for block in existing_credentials_blocks],
+            data=[
+                {"name": block.name}
+                for block in existing_credentials_blocks
+                if block.name is not None
+            ],
             opt_out_message="Create a new credentials block",
         )
 
@@ -766,6 +918,8 @@ async def prompt_select_blob_storage_credentials(
     credentials_block_schema = await client.get_most_recent_block_schema_for_block_type(
         block_type_id=credentials_block_type.id
     )
+    if credentials_block_schema is None:
+        raise ValueError(f"No schema found for {pretty_creds_block_type} block")
 
     console.print(
         f"\nProvide details on your new {pretty_storage_provider} credentials:"
@@ -774,7 +928,7 @@ async def prompt_select_blob_storage_credentials(
     hydrated_fields = {
         field_name: prompt(f"{field_name} [yellow]({props.get('type')})[/]")
         for field_name, props in credentials_block_schema.fields.get(
-            "properties"
+            "properties", {}
         ).items()
         if field_name in REQUIRED_FIELDS_FOR_CREDS_BLOCK[creds_block_type_slug]
     }

{prefect-3.1.11.dev2 → prefect-3.1.11.dev3}/src/prefect/cli/_types.py

@@ -5,7 +5,7 @@ Custom Prefect CLI types
 import asyncio
 import functools
 import sys
-from typing import Callable, List, Optional
+from typing import Any, Callable, List, Optional
 
 import typer
 from rich.console import Console
@@ -17,7 +17,7 @@ from prefect.settings import PREFECT_CLI_COLORS, Setting
 from prefect.utilities.asyncutils import is_async_fn
 
 
-def SettingsOption(setting: Setting, *args, **kwargs) -> typer.Option:
+def SettingsOption(setting: Setting, *args: Any, **kwargs: Any) -> Any:
     """Custom `typer.Option` factory to load the default value from settings"""
 
     return typer.Option(
@@ -33,7 +33,7 @@ def SettingsOption(setting: Setting, *args, **kwargs) -> typer.Option:
     )
 
 
-def SettingsArgument(setting: Setting, *args, **kwargs) -> typer.Argument:
+def SettingsArgument(setting: Setting, *args: Any, **kwargs: Any) -> Any:
     """Custom `typer.Argument` factory to load the default value from settings"""
 
     # See comments in `SettingsOption`
@@ -45,10 +45,12 @@ def SettingsArgument(setting: Setting, *args, **kwargs) -> typer.Argument:
     )
 
 
-def with_deprecated_message(warning: str):
-    def decorator(fn):
+def with_deprecated_message(
+    warning: str,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
         @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
             print("WARNING:", warning, file=sys.stderr, flush=True)
             return fn(*args, **kwargs)
 
@@ -66,12 +68,12 @@ class PrefectTyper(typer.Typer):
 
     def __init__(
         self,
-        *args,
+        *args: Any,
         deprecated: bool = False,
         deprecated_start_date: Optional[str] = None,
         deprecated_help: str = "",
         deprecated_name: str = "",
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(*args, **kwargs)
 
@@ -79,7 +81,7 @@ class PrefectTyper(typer.Typer):
         if self.deprecated:
             if not deprecated_name:
                 raise ValueError("Provide the name of the deprecated command group.")
-            self.deprecated_message = generate_deprecation_message(
+            self.deprecated_message: str = generate_deprecation_message(
                 name=f"The {deprecated_name!r} command group",
                 start_date=deprecated_start_date,
                 help=deprecated_help,
@@ -94,10 +96,10 @@ class PrefectTyper(typer.Typer):
     def add_typer(
         self,
         typer_instance: "PrefectTyper",
-        *args,
+        *args: Any,
         no_args_is_help: bool = True,
-        aliases: Optional[List[str]] = None,
-        **kwargs,
+        aliases: Optional[list[str]] = None,
+        **kwargs: Any,
     ) -> None:
         """
         This will cause help to be default command for all sub apps unless specifically stated otherwise, opposite of before.
@@ -120,14 +122,14 @@ class PrefectTyper(typer.Typer):
     def command(
         self,
         name: Optional[str] = None,
-        *args,
+        *args: Any,
         aliases: Optional[List[str]] = None,
         deprecated: bool = False,
         deprecated_start_date: Optional[str] = None,
         deprecated_help: str = "",
         deprecated_name: str = "",
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
         """
         Create a new command. If aliases are provided, the same command function
         will be registered with multiple names.
@@ -136,7 +138,7 @@ class PrefectTyper(typer.Typer):
         `deprecated_name` and `deprecated_start_date` must be provided.
         """
 
-        def wrapper(original_fn: Callable):
+        def wrapper(original_fn: Callable[..., Any]) -> Callable[..., Any]:
             # click doesn't support async functions, so we wrap them in
             # asyncio.run(). This has the advantage of keeping the function in
             # the main thread, which means signal handling works for e.g. the
@@ -149,10 +151,10 @@ class PrefectTyper(typer.Typer):
                 async_fn = original_fn
 
                 @functools.wraps(original_fn)
-                def sync_fn(*args, **kwargs):
+                def sync_fn(*args: Any, **kwargs: Any) -> Any:
                     return asyncio.run(async_fn(*args, **kwargs))
 
-                sync_fn.aio = async_fn
+                setattr(sync_fn, "aio", async_fn)
                 wrapped_fn = sync_fn
             else:
                 wrapped_fn = original_fn
@@ -195,7 +197,7 @@ class PrefectTyper(typer.Typer):
 
         return wrapper
 
-    def setup_console(self, soft_wrap: bool, prompt: bool):
+    def setup_console(self, soft_wrap: bool, prompt: bool) -> None:
         self.console = Console(
             highlight=False,
             color_system="auto" if PREFECT_CLI_COLORS else None,