diff --git a/airflow-core/src/airflow/utils/types.py b/airflow-core/src/airflow/utils/types.py index 28822919e382e..de50f46eb7e17 100644 --- a/airflow-core/src/airflow/utils/types.py +++ b/airflow-core/src/airflow/utils/types.py @@ -17,7 +17,7 @@ from __future__ import annotations import enum -from typing import TYPE_CHECKING, TypedDict +from typing import TYPE_CHECKING import airflow.sdk.definitions._internal.types @@ -57,12 +57,6 @@ def from_run_id(run_id: str) -> DagRunType: return DagRunType.MANUAL -class EdgeInfoType(TypedDict): - """Extra metadata that the DAG can store about an edge, usually generated from an EdgeModifier.""" - - label: str | None - - class DagRunTriggeredByType(enum.Enum): """Class with TriggeredBy types for DagRun.""" diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index 11c419cef04d4..555207444db2e 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -76,10 +76,9 @@ from airflow.sdk.definitions.abstractoperator import Operator from airflow.sdk.definitions.decorators import TaskDecoratorCollection + from airflow.sdk.definitions.edges import EdgeInfoType from airflow.sdk.definitions.taskgroup import TaskGroup from airflow.typing_compat import Self - from airflow.utils.types import EdgeInfoType - log = logging.getLogger(__name__) diff --git a/task-sdk/src/airflow/sdk/definitions/edges.py b/task-sdk/src/airflow/sdk/definitions/edges.py index 4bc620645dc4a..39fafc4b932c8 100644 --- a/task-sdk/src/airflow/sdk/definitions/edges.py +++ b/task-sdk/src/airflow/sdk/definitions/edges.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypedDict from airflow.sdk.definitions._internal.mixins import DependencyMixin @@ -187,3 +187,9 @@ def add_edge_info(self, dag: DAG, upstream_id: str, downstream_id: str): def Label(label: str): """Create an EdgeModifier that sets a human-readable label on the edge.""" return EdgeModifier(label=label) + + +class EdgeInfoType(TypedDict): + """Extra metadata that the DAG can store about an edge, usually generated from an EdgeModifier.""" + + label: str | None