This is an automated email from the ASF dual-hosted git repository.
amoghdesai pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e44020f5310 Feature: enhance task-sdk definitions context stub file
typing (#54421)
e44020f5310 is described below
commit e44020f531023f48737157f0d274626cc1dec084
Author: Nataneljpwd <[email protected]>
AuthorDate: Fri Jan 9 11:45:29 2026 +0000
Feature: enhance task-sdk definitions context stub file typing (#54421)
---
task-sdk/src/airflow/sdk/definitions/context.py | 26 +++++++++++++------------
1 file changed, 14 insertions(+), 12 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/definitions/context.py
b/task-sdk/src/airflow/sdk/definitions/context.py
index 46daccb6f56..ba4d9b65905 100644
--- a/task-sdk/src/airflow/sdk/definitions/context.py
+++ b/task-sdk/src/airflow/sdk/definitions/context.py
@@ -22,6 +22,8 @@ import os
from collections.abc import MutableMapping
from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, cast
+from typing_extensions import NotRequired
+
if TYPE_CHECKING:
import jinja2
from pendulum import DateTime
@@ -43,25 +45,25 @@ class Context(TypedDict, total=False):
conn: Any
dag: DAG
dag_run: DagRunProtocol
- data_interval_end: DateTime | None
- data_interval_start: DateTime | None
+ data_interval_end: NotRequired[DateTime | None]
+ data_interval_start: NotRequired[DateTime | None]
outlet_events: OutletEventAccessorsProtocol
ds: str
ds_nodash: str
- expanded_ti_count: int | None
- exception: None | str | BaseException
+ expanded_ti_count: NotRequired[int | None]
+ exception: NotRequired[None | str | BaseException]
inlets: list
inlet_events: InletEventsAccessors
logical_date: DateTime
macros: Any
- map_index_template: str | None
+ map_index_template: NotRequired[str | None]
outlets: list
params: dict[str, Any]
- prev_data_interval_start_success: DateTime | None
- prev_data_interval_end_success: DateTime | None
- prev_start_date_success: DateTime | None
- prev_end_date_success: DateTime | None
- reason: str | None
+ prev_data_interval_start_success: NotRequired[DateTime | None]
+ prev_data_interval_end_success: NotRequired[DateTime | None]
+ prev_start_date_success: NotRequired[DateTime | None]
+ prev_end_date_success: NotRequired[DateTime | None]
+ reason: NotRequired[str | None]
run_id: str
start_date: DateTime
# TODO: Remove Operator from below once we have MappedOperator to the Task
SDK
@@ -71,12 +73,12 @@ class Context(TypedDict, total=False):
task_instance: RuntimeTaskInstanceProtocol
task_instance_key_str: str
# `templates_dict` is only set in PythonOperator
- templates_dict: dict[str, Any] | None
+ templates_dict: NotRequired[dict[str, Any] | None]
test_mode: bool
ti: RuntimeTaskInstanceProtocol
# triggering_asset_events: Mapping[str, Collection[AssetEvent |
AssetEventPydantic]]
triggering_asset_events: Any
- try_number: int | None
+ try_number: NotRequired[int | None]
ts: str
ts_nodash: str
ts_nodash_with_tz: str
