Add static check to keep @task decorator type hints in sync with their operators

[sjyangkevin]

Several Airflow providers offer @task.<name> decorators, such as @task.python, @task.kubernetes, @task.docker and @task.llm. Each decorator builds an operator behind the scenes, so the options it accepts are defined in that operator's code. The type hints that editors and type checkers show for these decorators come from somewhere else: a typing file in the Task SDK, task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi. No hook currently checks that the two agree, so as operators and decorators evolve, the typing file can fall out of step without anyone noticing.

This PR adds a static check that compares the typing file with the source code of every @task.<name> decorator registered in the repository, and brings the typing file up to date.

closes: #48448
related: #51048

Current state

Keeping the typing file in sync seems to be a manual step today. As the issue notes, a step like this is easy to overlook, especially since it applies to decorators spread across several providers. Comparing the typing file with the current source shows where it has fallen behind:

• 12 registered decorators don't have a typing entry yet: agent, analytics, llm, llm_batch, llm_branch, llm_file_analysis, llm_schema_compare, llm_sql, sftp_sensor, snowpark, sql and stub. Editors can't suggest their options, and type checkers don't check them.
• 13 entries don't yet include every option their operators accept. For example, the two Kubernetes entries don't include about a dozen KubernetesPodOperator options, such as callbacks, logging_interval and runtime_class_name.
• A few entries list options the decorators don't accept. @task.docker lists use_dill, which that decorator doesn't take. @task.llm_schema_compare lists output_type, which the operator rejects with a TypeError, and serialize_output, whose value the operator always sets itself. In these cases a type checker accepts code that fails, or has no effect, at runtime.

What this PR changes

1. Adds a prek hook, check-task-decorator-arguments, that fails whenever a decorator's entry in the typing file doesn't match the options the decorator actually accepts.
2. Brings the typing file up to date: adds entries for the 12 decorators that don't have one, and updates 13 existing entries.
3. Updates the reminder comment in KubernetesPodOperator to point at the typing file.

How the check works

The check needs no list of decorators of its own. For each decorator it follows the same trail a reader would follow by hand:

1. Find the decorators. It reads the task-decorators section of every provider.yaml, which is the same list Airflow uses at runtime to find @task.<name> decorators. A new provider or decorator is picked up automatically.

2. Find the operator each decorator builds. It opens the registered decorator function and finds the class passed as decorated_operator_class. That is the operator created for every task the decorator produces.

3. Follow the operator's parent classes. It resolves each parent class by reading the module's imports, including the compat modules that re-export classes. This works across providers: the Snowflake decorator's operator, for example, inherits from the standard provider's PythonOperator. It puts the classes in the order Python itself would call their constructors, and stops at BaseOperator. BaseOperator passes nothing further and rejects leftover arguments, so no class after it can ever receive a user option.

4. Collect every accepted option. Each constructor along that path contributes its parameters, as do the decorator function and Airflow's shared decorator factory.

5. Drop the options the decorator fills in itself. Some options are always set by the decorator, so users can't pass them:

   • Kubernetes always sets cmds.
   • Bash, SQL and the LLM decorators set bash_command, sql and prompt from the value the decorated function returns.
   • LLMSchemaCompareOperator always forces serialize_output=True.

   The check recognises these because the value passed on is fixed, rather than taken from the user's input. A value read from the user's arguments, such as name=kwargs.pop("name", default), stays a user option.

6. Compare with the typing entry. It reports options that are missing from the entry and options the entry lists but the decorator doesn't accept. Each message names the file and line where the option is defined or filled in.

Everything is read as source code with Python's ast module. Nothing is imported, so the check runs on a clean checkout, without provider dependencies or Breeze, in about 2 seconds. Example output:

__init__.pyi:542: @task.kubernetes is missing parameters: logging_interval. Declared in providers/cncf/kubernetes/.../operators/pod.py:325.
__init__.pyi:382: @task.docker declares unexpected parameters: command. command is internally supplied at providers/docker/.../decorators/docker.py:125.

When the check fails

Situation	What the check reports	What the contributor does
A new provider ships a decorator, or a provider adds one	The decorator has no typing entry, or its entry is missing options, each with its source line	Add the entry
A decorator is removed from provider.yaml	The typing file still has an entry for it	Delete the entry
A decorator is renamed	The new name needs an entry, and the old entry is left over	Rename the entry
provider.yaml points at a missing file or function	The missing file or function	Fix the registration
An operator gains, loses or renames an option	Every decorator built on it, including decorators in other providers (a PythonOperator change affects 9 decorators across 3 providers)	Update those entries
An operator gets a new parent class	The new parent's options are picked up automatically	Update the entries
A decorator starts filling an option itself	The option is now unexpected, with the line where it is filled	Remove it from the entry
Code is written in a way the check can't follow, e.g. picking the operator class at runtime	An explanation of what it couldn't resolve. It never silently passes	Use the usual pattern

None of these cases require editing the check itself. Contributors only touch the typing file, and the message says exactly what to change. The check also requires plain @task to accept the same options as @task.python.

Example runs: real output, one edit per scenario

Each scenario starts from this branch, makes one edit, and runs prek run check-task-decorator-arguments --all-files. Long paths and option lists are shortened with ….

On this branch as-is:

Check registered task decorator coverage and argument declarations....Passed

1. A provider registers a new decorator without a typing entry. This scenario uses the provider.yaml entry and decorator module from open PR #72033 (@task.eks_pod):

Check registered task decorator coverage and argument declarations....Failed
- hook id: check-task-decorator-arguments
- exit code: 1

  @task.eks_pod: registered without a named typing declaration. Update task-sdk/…/decorators/__init__.pyi.
  @task.eks_pod: task-sdk/…/decorators/__init__.pyi: expected one configuration overload TaskDecoratorCollection.eks_pod, found 0

After an empty eks_pod entry is added, the check lists every option to declare and where each one comes from, here across two providers:

  …/__init__.pyi:63: @task.eks_pod is missing parameters: active_deadline_seconds, affinity, annotations, …, aws_conn_id, …, cluster_name, …, region, …, xcom_sidecar_container_security_context. Declared in providers/amazon/…/decorators/eks.py:72, providers/amazon/…/operators/eks.py:1136, providers/cncf/kubernetes/…/operators/pod.py:325. Update …/__init__.pyi.

2. An operator gains an option. Adding cache_result: bool = False to PythonOperator.__init__ flags every decorator built on it, across three providers:

  …/__init__.pyi:226: @task.branch is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:310: @task.branch_external_python is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:246: @task.branch_virtualenv is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:177: @task.external_python is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:962: @task.pyspark is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:64: @task.python is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:360: @task.short_circuit is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:1032: @task.snowpark is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
  …/__init__.pyi:107: @task.virtualenv is missing parameters: cache_result. Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.

3. A typing entry lists an option the decorator always fills itself. Adding cmds to the @task.kubernetes entry:

  …/__init__.pyi:542: @task.kubernetes declares unexpected parameters: cmds. cmds is internally supplied at providers/cncf/kubernetes/…/decorators/kubernetes.py:72. Update …/__init__.pyi to match providers/cncf/kubernetes/…/decorators/kubernetes.py:_KubernetesDecoratedOperator and the constructors it inherits.

4. A decorator is removed. Deleting the task-decorators block from the Docker provider's provider.yaml:

  @task.docker: typing declaration has no provider registration. Update task-sdk/…/__init__.pyi.

5. Code the check can't follow. Choosing the operator class at runtime with decorated_operator_class=_get_operator_class():

  @task.docker: providers/docker/…/decorators/docker.py:213: cannot resolve decorated_operator_class _get_operator_class()

When the check deliberately passes:

6. A new option on BaseOperator. Adding new_shared_option to BaseOperator.__init__ → Passed. Shared options go through each entry's **kwargs, so the 25 entries don't need editing.

7. A new operator option that the decorator fills itself. Adding script_path to BashOperator, with _BashDecoratedOperator passing script_path=SET_DURING_EXECUTION → Passed. Users can't set the option, so the entry must not list it. Without the decorator's line, the same edit fails with @task.bash is missing parameters: script_path.

8. Changing an option's type or default. Changing image: str | None = None to image: str = "busybox" in KubernetesPodOperator → Passed. Only option names are compared; see below.

What the check does not cover

• Option types and default values. Only option names are compared. A few entries on main have types narrower than their operators. For example, @task.sensor declares poke_interval: float although timedelta works. Those can be fixed in a follow-up.
• Options shared by every operator (task_id, retries, pool, …). The typing entries deliberately accept these through their trailing **kwargs instead of listing about 55 options in every entry. The check keeps that convention: it stops at BaseOperator, and it reports a shared option listed in a single entry as unexpected.
• @task.run_if and @task.skip_if. They wrap another task rather than creating an operator, and their signatures are defined in the Task SDK itself.
• Decorators from providers outside this repository. They still work at runtime, but this repository's typing file can't describe them.
• Options an operator reads from **kwargs without declaring them. This happens in 3 operator classes across the providers (two Databricks sensors and one Cloud Composer sensor). None of them has a decorator today.
• Other decorator families, such as @dag, @task_group and @asset.

Impact on CI and pre-commit

• CI: the static checks job already runs prek --all-files, so this is one more hook taking about 2 seconds. It needs no Breeze image and imports no provider code.
• Local commits: the hook uses always_run: true, so it runs on every commit, adding about 2 seconds. This is needed because deleted files are not passed to hooks, and removing a decorator's files would otherwise go unnoticed. The check always scans all decorators, because a change to one shared file, such as PythonOperator, the decorator factory or the typing file itself, can affect decorators in several providers.
• When it can block a PR: when that PR leaves the typing file out of sync, which is the intended case. It can also block a PR when two open PRs each pass on their own but conflict once both are merged. For example, open PR Add @task.eks_pod TaskFlow decorator for EksPodOperator #72033 registers a new @task.eks_pod decorator without a typing entry, so whichever of the two PRs merges second will need to add that entry. The error message lists exactly what to add.
• Coverage today: all 25 @task.<name> decorators across 8 providers. Every module that builds a decorator with the shared factory is registered in a provider.yaml, so nothing is left unchecked.

Question for reviewers: newsfragment?

The typing changes alter what type checkers accept for @task.<name> calls. Twelve decorators now have typed signatures, and options the decorators don't accept were removed. Runtime behaviour does not change. Since the Task SDK ships with airflow-core, would you like an airflow-core/newsfragments entry for this, and if so of which type? I haven't added one, following the guideline to add newsfragments only when the change is clearly user-facing.

Testing

• uv run --project scripts pytest scripts/tests/: 1,814 passed, 1 skipped. This includes 133 tests for the new hook. Each rule in the check was also disabled in turn to confirm that a test fails (13 of 13).
• prek run --from-ref <base> --stage pre-commit: passed.
• prek run --from-ref <base> --stage manual --skip upgrade-important-versions: passed, except mypy-providers, which couldn't run because my local Breeze CI image is out of date. The only provider file it would check is pod.py, where the only change is a comment. CI runs it.
• prek run mypy-task-sdk --all-files: passed. This type-checks the typing file.
• breeze build-docs task-sdk: passed. The Task SDK API docs are generated from the typing file.
• Ran the derivation over all 1,262 operator and sensor classes in the providers: every one resolved through to BaseOperator.
• Not run locally: the core and provider test suites breeze verify selects (Kubernetes, Amazon, Google, …). They are selected because of the .pyi file and a comment-only change in pod.py, neither of which changes runtime behaviour. CI runs them.

---

Was generative AI tooling used to co-author this PR?

• Yes — Claude Code (Opus 5.5)

Generated-by: Claude Code (Opus 5.5) following the guidelines

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.