You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "sjyangkevin (via GitHub)" <gi...@apache.org> on 2026/10/14 03:55:40 UTC
[PR] Add static check to keep @task decorator type hints in sync with their operators [airflow]
sjyangkevin opened a new pull request, #74689:
URL: https://github.com/apache/airflow/pull/74689
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
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`.
<details>
<summary><b>Example runs</b>: real output, one edit per scenario</summary>
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.
</details>
## 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 #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?
- [X] Yes — Claude Code (Opus 5.5)
Generated-by: Claude Code (Opus 5.5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add static check to keep @task decorator type hints in sync with their operators [airflow]
Posted by "sjyangkevin (via GitHub)" <gi...@apache.org>.
sjyangkevin commented on PR #74689:
URL: https://github.com/apache/airflow/pull/74689#issuecomment-6160542438
Note: this check would fail the current open PR https://github.com/apache/airflow/pull/72033 as shown in the **Example run** section. A decorator is defined, but without updating the task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org