Skip to content

Commit

Permalink
Raise type completeness to 99.8% (#16734)
Browse files Browse the repository at this point in the history
desertaxle authored Jan 16, 2025

Verified

This commit was signed with the committer’s verified signature.
frostming Frost Ming
1 parent f8255b1 commit ccbf65f
Showing 56 changed files with 441 additions and 295 deletions.
10 changes: 6 additions & 4 deletions src/prefect/_internal/schemas/validators.py
Original file line number Diff line number Diff line change
@@ -6,6 +6,8 @@
This will be subject to consolidation and refactoring over the next few months.
"""

from __future__ import annotations

import os
import re
import urllib.parse
@@ -627,18 +629,18 @@ def validate_name_present_on_nonanonymous_blocks(values: M) -> M:


@overload
def validate_command(v: str) -> Path:
def validate_working_dir(v: str) -> Path:
...


@overload
def validate_command(v: None) -> None:
def validate_working_dir(v: None) -> None:
...


def validate_command(v: Optional[str]) -> Optional[Path]:
def validate_working_dir(v: Optional[Path | str]) -> Optional[Path]:
"""Make sure that the working directory is formatted for the current platform."""
if v is not None:
if isinstance(v, str):
return relative_path_to_current_platform(v)
return v

5 changes: 4 additions & 1 deletion src/prefect/artifacts.py
Original file line number Diff line number Diff line change
@@ -20,7 +20,10 @@
from prefect.utilities.asyncutils import sync_compatible
from prefect.utilities.context import get_task_and_flow_run_ids

logger = get_logger("artifacts")
if TYPE_CHECKING:
import logging

logger: "logging.Logger" = get_logger("artifacts")

if TYPE_CHECKING:
from prefect.client.orchestration import PrefectClient
2 changes: 1 addition & 1 deletion src/prefect/automations.py
Original file line number Diff line number Diff line change
@@ -137,7 +137,7 @@ def create(self: Self) -> Self:
self.id = client.create_automation(automation=automation)
return self

async def aupdate(self: Self):
async def aupdate(self: Self) -> None:
"""
Updates an existing automation.
4 changes: 2 additions & 2 deletions src/prefect/cli/dashboard.py
Original file line number Diff line number Diff line change
@@ -8,15 +8,15 @@
from prefect.settings import PREFECT_UI_URL
from prefect.utilities.asyncutils import run_sync_in_worker_thread

dashboard_app = PrefectTyper(
dashboard_app: PrefectTyper = PrefectTyper(
name="dashboard",
help="Commands for interacting with the Prefect UI.",
)
app.add_typer(dashboard_app)


@dashboard_app.command()
async def open():
async def open() -> None:
"""
Open the Prefect UI in the browser.
"""
4 changes: 2 additions & 2 deletions src/prefect/cli/dev.py
Original file line number Diff line number Diff line change
@@ -35,13 +35,13 @@
Note that many of these commands require extra dependencies (such as npm and MkDocs)
to function properly.
"""
dev_app = PrefectTyper(
dev_app: PrefectTyper = PrefectTyper(
name="dev", short_help="Internal Prefect development.", help=DEV_HELP
)
app.add_typer(dev_app)


def exit_with_error_if_not_editable_install():
def exit_with_error_if_not_editable_install() -> None:
if (
prefect.__module_path__.parent == "site-packages"
or not (prefect.__development_base_path__ / "setup.py").exists()
6 changes: 3 additions & 3 deletions src/prefect/cli/events.py
Original file line number Diff line number Diff line change
@@ -15,7 +15,7 @@
get_events_subscriber,
)

events_app = PrefectTyper(name="events", help="Stream events.")
events_app: PrefectTyper = PrefectTyper(name="events", help="Stream events.")
app.add_typer(events_app, aliases=["event"])


@@ -60,7 +60,7 @@ async def stream(
handle_error(exc)


async def handle_event(event: Event, format: StreamFormat, output_file: str):
async def handle_event(event: Event, format: StreamFormat, output_file: str) -> None:
if format == StreamFormat.json:
event_data = orjson.dumps(event.model_dump(), default=str).decode()
elif format == StreamFormat.text:
@@ -74,7 +74,7 @@ async def handle_event(event: Event, format: StreamFormat, output_file: str):
print(event_data)


def handle_error(exc):
def handle_error(exc: Exception) -> None:
if isinstance(exc, websockets.exceptions.ConnectionClosedError):
exit_with_error(f"Connection closed, retrying... ({exc})")
elif isinstance(exc, (KeyboardInterrupt, asyncio.exceptions.CancelledError)):
2 changes: 1 addition & 1 deletion src/prefect/cli/flow.py
Original file line number Diff line number Diff line change
@@ -19,7 +19,7 @@
from prefect.runner import Runner
from prefect.utilities import urls

flow_app = PrefectTyper(name="flow", help="View and serve flows.")
flow_app: PrefectTyper = PrefectTyper(name="flow", help="View and serve flows.")
app.add_typer(flow_app, aliases=["flows"])


6 changes: 4 additions & 2 deletions src/prefect/cli/flow_run.py
Original file line number Diff line number Diff line change
@@ -28,13 +28,15 @@
from prefect.runner import Runner
from prefect.states import State

flow_run_app = PrefectTyper(name="flow-run", help="Interact with flow runs.")
flow_run_app: PrefectTyper = PrefectTyper(
name="flow-run", help="Interact with flow runs."
)
app.add_typer(flow_run_app, aliases=["flow-runs"])

LOGS_DEFAULT_PAGE_SIZE = 200
LOGS_WITH_LIMIT_FLAG_DEFAULT_NUM_LOGS = 20

logger = get_logger(__name__)
logger: "logging.Logger" = get_logger(__name__)


@flow_run_app.command()
2 changes: 1 addition & 1 deletion src/prefect/cli/global_concurrency_limit.py
Original file line number Diff line number Diff line change
@@ -22,7 +22,7 @@
PrefectHTTPStatusError,
)

global_concurrency_limit_app = PrefectTyper(
global_concurrency_limit_app: PrefectTyper = PrefectTyper(
name="global-concurrency-limit",
help="Manage global concurrency limits.",
)
10 changes: 6 additions & 4 deletions src/prefect/cli/profile.py
Original file line number Diff line number Diff line change
@@ -26,7 +26,9 @@
from prefect.settings import ProfilesCollection
from prefect.utilities.collections import AutoEnum

profile_app = PrefectTyper(name="profile", help="Select and manage Prefect profiles.")
profile_app: PrefectTyper = PrefectTyper(
name="profile", help="Select and manage Prefect profiles."
)
app.add_typer(profile_app, aliases=["profiles"])

_OLD_MINIMAL_DEFAULT_PROFILE_CONTENT: str = """active = "default"
@@ -263,8 +265,8 @@ def inspect(

def show_profile_changes(
user_profiles: ProfilesCollection, default_profiles: ProfilesCollection
):
changes = []
) -> bool:
changes: list[tuple[str, str]] = []

for name in default_profiles.names:
if name not in user_profiles:
@@ -343,7 +345,7 @@ class ConnectionStatus(AutoEnum):
INVALID_API = AutoEnum.auto()


async def check_server_connection():
async def check_server_connection() -> ConnectionStatus:
httpx_settings = dict(timeout=3)
try:
# attempt to infer Cloud 2.0 API from the connection URL
8 changes: 4 additions & 4 deletions src/prefect/cli/root.py
Original file line number Diff line number Diff line change
@@ -26,16 +26,16 @@
PREFECT_TEST_MODE,
)

app = PrefectTyper(add_completion=True, no_args_is_help=True)
app: PrefectTyper = PrefectTyper(add_completion=True, no_args_is_help=True)


def version_callback(value: bool):
def version_callback(value: bool) -> None:
if value:
print(prefect.__version__)
raise typer.Exit()


def is_interactive():
def is_interactive() -> bool:
return app.console.is_interactive


@@ -157,7 +157,7 @@ def get_prefect_integrations() -> Dict[str, str]:
return integrations


def display(object: Dict[str, Any], nesting: int = 0):
def display(object: Dict[str, Any], nesting: int = 0) -> None:
"""Recursive display of a dictionary with nesting."""
for key, value in object.items():
key += ":"
20 changes: 14 additions & 6 deletions src/prefect/cli/server.py
Original file line number Diff line number Diff line change
@@ -15,6 +15,7 @@
import textwrap
from pathlib import Path
from types import ModuleType
from typing import TYPE_CHECKING

import typer
import uvicorn
@@ -48,23 +49,30 @@
from prefect.settings.context import temporary_settings
from prefect.utilities.asyncutils import run_sync_in_worker_thread

server_app = PrefectTyper(
if TYPE_CHECKING:
import logging

server_app: PrefectTyper = PrefectTyper(
name="server",
help="Start a Prefect server instance and interact with the database",
)
database_app = PrefectTyper(name="database", help="Interact with the database.")
services_app = PrefectTyper(name="services", help="Interact with server loop services.")
database_app: PrefectTyper = PrefectTyper(
name="database", help="Interact with the database."
)
services_app: PrefectTyper = PrefectTyper(
name="services", help="Interact with server loop services."
)
server_app.add_typer(database_app)
server_app.add_typer(services_app)
app.add_typer(server_app)

logger = get_logger(__name__)
logger: "logging.Logger" = get_logger(__name__)

SERVER_PID_FILE_NAME = "server.pid"
SERVICES_PID_FILE = Path(PREFECT_HOME.value()) / "services.pid"


def generate_welcome_blurb(base_url: str, ui_enabled: bool):
def generate_welcome_blurb(base_url: str, ui_enabled: bool) -> str:
blurb = textwrap.dedent(
r"""
___ ___ ___ ___ ___ ___ _____
@@ -109,7 +117,7 @@ def generate_welcome_blurb(base_url: str, ui_enabled: bool):
return blurb


def prestart_check(base_url: str):
def prestart_check(base_url: str) -> None:
"""
Check if `PREFECT_API_URL` is set in the current profile. If not, prompt the user to set it.
8 changes: 4 additions & 4 deletions src/prefect/cli/shell.py
Original file line number Diff line number Diff line change
@@ -8,7 +8,7 @@
import subprocess
import sys
import threading
from typing import Any, Dict, List, Optional
from typing import IO, Any, Callable, Dict, List, Optional

import typer
from typing_extensions import Annotated
@@ -25,13 +25,13 @@
from prefect.settings import PREFECT_UI_URL
from prefect.types.entrypoint import EntrypointType

shell_app = PrefectTyper(
shell_app: PrefectTyper = PrefectTyper(
name="shell", help="Serve and watch shell commands as Prefect flows."
)
app.add_typer(shell_app)


def output_stream(pipe, logger_function):
def output_stream(pipe: IO[str], logger_function: Callable[[str], None]) -> None:
"""
Read from a pipe line by line and log using the provided logging function.
@@ -44,7 +44,7 @@ def output_stream(pipe, logger_function):
logger_function(line.strip())


def output_collect(pipe, container):
def output_collect(pipe: IO[str], container: list[str]) -> None:
"""
Collects output from a subprocess pipe and stores it in a container list.
2 changes: 1 addition & 1 deletion src/prefect/cli/task.py
Original file line number Diff line number Diff line change
@@ -8,7 +8,7 @@
from prefect.task_worker import serve as task_serve
from prefect.utilities.importtools import import_object

task_app = PrefectTyper(name="task", help="Work with task scheduling.")
task_app: PrefectTyper = PrefectTyper(name="task", help="Work with task scheduling.")
app.add_typer(task_app, aliases=["task"])


4 changes: 3 additions & 1 deletion src/prefect/cli/task_run.py
Original file line number Diff line number Diff line change
@@ -22,7 +22,9 @@
from prefect.client.schemas.sorting import LogSort, TaskRunSort
from prefect.exceptions import ObjectNotFound

task_run_app = PrefectTyper(name="task-run", help="View and inspect task runs.")
task_run_app: PrefectTyper = PrefectTyper(
name="task-run", help="View and inspect task runs."
)
app.add_typer(task_run_app, aliases=["task-runs"])

LOGS_DEFAULT_PAGE_SIZE = 200
2 changes: 1 addition & 1 deletion src/prefect/cli/variable.py
Original file line number Diff line number Diff line change
@@ -14,7 +14,7 @@
from prefect.client.schemas.actions import VariableCreate, VariableUpdate
from prefect.exceptions import ObjectNotFound

variable_app = PrefectTyper(name="variable", help="Manage variables.")
variable_app: PrefectTyper = PrefectTyper(name="variable", help="Manage variables.")
app.add_typer(variable_app)


4 changes: 2 additions & 2 deletions src/prefect/cli/work_pool.py
Original file line number Diff line number Diff line change
@@ -32,11 +32,11 @@
get_default_base_job_template_for_infrastructure_type,
)

work_pool_app = PrefectTyper(name="work-pool", help="Manage work pools.")
work_pool_app: PrefectTyper = PrefectTyper(name="work-pool", help="Manage work pools.")
app.add_typer(work_pool_app, aliases=["work-pool"])


def set_work_pool_as_default(name: str):
def set_work_pool_as_default(name: str) -> None:
profile = update_current_profile({"PREFECT_DEFAULT_WORK_POOL_NAME": name})
app.console.print(
f"Set {name!r} as default work pool for profile {profile.name!r}\n",
2 changes: 1 addition & 1 deletion src/prefect/cli/work_queue.py
Original file line number Diff line number Diff line change
@@ -20,7 +20,7 @@
from prefect.client.schemas.objects import DEFAULT_AGENT_WORK_POOL_NAME
from prefect.exceptions import ObjectAlreadyExists, ObjectNotFound

work_app = PrefectTyper(name="work-queue", help="Manage work queues.")
work_app: PrefectTyper = PrefectTyper(name="work-queue", help="Manage work queues.")
app.add_typer(work_app, aliases=["work-queues"])


4 changes: 3 additions & 1 deletion src/prefect/cli/worker.py
Original file line number Diff line number Diff line change
@@ -28,7 +28,9 @@
)
from prefect.workers.base import BaseWorker

worker_app = PrefectTyper(name="worker", help="Start and interact with workers.")
worker_app: PrefectTyper = PrefectTyper(
name="worker", help="Start and interact with workers."
)
app.add_typer(worker_app)


11 changes: 7 additions & 4 deletions src/prefect/deployments/base.py
Original file line number Diff line number Diff line change
@@ -19,6 +19,7 @@
from prefect.client.schemas.objects import ConcurrencyLimitStrategy
from prefect.client.schemas.schedules import IntervalSchedule
from prefect.utilities._git import get_git_branch, get_git_remote_origin_url
from prefect.utilities.annotations import NotSet
from prefect.utilities.filesystem import create_default_ignore_file
from prefect.utilities.templating import apply_values

@@ -113,7 +114,9 @@ def create_default_prefect_yaml(
return True


def configure_project_by_recipe(recipe: str, **formatting_kwargs) -> dict:
def configure_project_by_recipe(
recipe: str, **formatting_kwargs: Any
) -> dict[str, Any] | type[NotSet]:
"""
Given a recipe name, returns a dictionary representing base configuration options.
@@ -131,13 +134,13 @@ def configure_project_by_recipe(recipe: str, **formatting_kwargs) -> dict:
raise ValueError(f"Unknown recipe {recipe!r} provided.")

with recipe_path.open(mode="r") as f:
config = yaml.safe_load(f)
config: dict[str, Any] = yaml.safe_load(f)

config = apply_values(
templated_config = apply_values(
template=config, values=formatting_kwargs, remove_notset=False
)

return config
return templated_config


def initialize_project(
Loading

0 comments on commit ccbf65f

Please sign in to comment.