This is an automated email from the ASF dual-hosted git repository.
kaxilnaik 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 fb67e128cf2 Move BaseOperator to `airflow/sdk/bases/operator.py`
(#48529)
fb67e128cf2 is described below
commit fb67e128cf2bb2ab7df24b0958ed45761c779e7f
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Mar 29 15:54:19 2025 +0530
Move BaseOperator to `airflow/sdk/bases/operator.py` (#48529)
---
airflow-core/src/airflow/decorators/base.py | 2 +-
airflow-core/src/airflow/decorators/condition.py | 2 +-
airflow-core/src/airflow/models/abstractoperator.py | 2 +-
airflow-core/src/airflow/models/baseoperator.py | 2 +-
airflow-core/src/airflow/models/taskinstance.py | 2 +-
airflow-core/src/airflow/models/taskmap.py | 2 +-
airflow-core/src/airflow/serialization/serialized_objects.py | 2 +-
airflow-core/src/airflow/utils/task_group.py | 2 +-
dev/mypy/plugin/outputs.py | 2 +-
devel-common/src/tests_common/pytest_plugin.py | 2 +-
.../airflow/providers/openlineage/utils/selective_enable.py | 2 +-
.../tests/unit/openlineage/extractors/test_manager.py | 2 +-
scripts/ci/pre_commit/check_base_operator_partial_arguments.py | 2 +-
task-sdk/src/airflow/sdk/__init__.py | 10 +++++-----
.../src/airflow/sdk/bases/{baseoperator.py => operator.py} | 0
task-sdk/src/airflow/sdk/bases/sensor.py | 2 +-
.../src/airflow/sdk/definitions/_internal/abstractoperator.py | 2 +-
task-sdk/src/airflow/sdk/definitions/_internal/mixins.py | 2 +-
task-sdk/src/airflow/sdk/definitions/_internal/node.py | 2 +-
task-sdk/src/airflow/sdk/definitions/context.py | 2 +-
task-sdk/src/airflow/sdk/definitions/dag.py | 2 +-
task-sdk/src/airflow/sdk/definitions/mappedoperator.py | 2 +-
task-sdk/src/airflow/sdk/definitions/taskgroup.py | 2 +-
task-sdk/src/airflow/sdk/definitions/xcom_arg.py | 2 +-
task-sdk/src/airflow/sdk/execution_time/context.py | 2 +-
task-sdk/src/airflow/sdk/execution_time/task_runner.py | 2 +-
task-sdk/src/airflow/sdk/types.py | 2 +-
task-sdk/tests/task_sdk/bases/test_baseoperator.py | 2 +-
task-sdk/tests/task_sdk/dags/super_basic.py | 2 +-
task-sdk/tests/task_sdk/dags/super_basic_run.py | 2 +-
task-sdk/tests/task_sdk/definitions/test_dag.py | 2 +-
task-sdk/tests/task_sdk/definitions/test_mappedoperator.py | 2 +-
task-sdk/tests/task_sdk/definitions/test_mixins.py | 2 +-
33 files changed, 36 insertions(+), 36 deletions(-)
diff --git a/airflow-core/src/airflow/decorators/base.py
b/airflow-core/src/airflow/decorators/base.py
index d0c8bdcf6da..c93f9fc2716 100644
--- a/airflow-core/src/airflow/decorators/base.py
+++ b/airflow-core/src/airflow/decorators/base.py
@@ -41,7 +41,7 @@ from airflow.models.expandinput import (
ListOfDictsExpandInput,
is_mappable,
)
-from airflow.sdk.bases.baseoperator import BaseOperator as TaskSDKBaseOperator
+from airflow.sdk.bases.operator import BaseOperator as TaskSDKBaseOperator
from airflow.sdk.definitions._internal.contextmanager import DagContext,
TaskGroupContext
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.definitions.mappedoperator import MappedOperator,
ensure_xcomarg_return_value
diff --git a/airflow-core/src/airflow/decorators/condition.py
b/airflow-core/src/airflow/decorators/condition.py
index a38b2ab2419..e685d768751 100644
--- a/airflow-core/src/airflow/decorators/condition.py
+++ b/airflow-core/src/airflow/decorators/condition.py
@@ -25,7 +25,7 @@ from airflow.exceptions import AirflowSkipException
if TYPE_CHECKING:
from typing_extensions import TypeAlias
- from airflow.sdk.bases.baseoperator import TaskPreExecuteHook
+ from airflow.sdk.bases.operator import TaskPreExecuteHook
from airflow.sdk.definitions.context import Context
BoolConditionFunc: TypeAlias = Callable[[Context], bool]
diff --git a/airflow-core/src/airflow/models/abstractoperator.py
b/airflow-core/src/airflow/models/abstractoperator.py
index 9dc8f32bdfe..0965d9550d6 100644
--- a/airflow-core/src/airflow/models/abstractoperator.py
+++ b/airflow-core/src/airflow/models/abstractoperator.py
@@ -33,7 +33,7 @@ from airflow.utils.weight_rule import db_safe_priority
if TYPE_CHECKING:
from sqlalchemy.orm import Session
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.task.priority_strategy import PriorityWeightStrategy
from airflow.triggers.base import StartTriggerArgs
diff --git a/airflow-core/src/airflow/models/baseoperator.py
b/airflow-core/src/airflow/models/baseoperator.py
index 23181f97d58..f052030f8df 100644
--- a/airflow-core/src/airflow/models/baseoperator.py
+++ b/airflow-core/src/airflow/models/baseoperator.py
@@ -47,7 +47,7 @@ from airflow.models.abstractoperator import (
NotMapped,
)
from airflow.models.taskinstance import TaskInstance, clear_task_instances
-from airflow.sdk.bases.baseoperator import (
+from airflow.sdk.bases.operator import (
BaseOperator as TaskSDKBaseOperator,
# Re-export for compat
chain as chain,
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index bb1d4016fdb..85d5f21d530 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -605,7 +605,7 @@ def _execute_task(task_instance: TaskInstance, context:
Context, task_orig: Oper
:meta private:
"""
- from airflow.sdk.bases.baseoperator import ExecutorSafeguard
+ from airflow.sdk.bases.operator import ExecutorSafeguard
from airflow.sdk.definitions.mappedoperator import MappedOperator
task_to_execute = task_instance.task
diff --git a/airflow-core/src/airflow/models/taskmap.py
b/airflow-core/src/airflow/models/taskmap.py
index 12f386c3b45..f0fd4c0231b 100644
--- a/airflow-core/src/airflow/models/taskmap.py
+++ b/airflow-core/src/airflow/models/taskmap.py
@@ -133,7 +133,7 @@ class TaskMap(TaskInstanceDependencies):
from airflow.models.baseoperator import BaseOperator as DBBaseOperator
from airflow.models.expandinput import NotFullyPopulated
from airflow.models.taskinstance import TaskInstance
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.settings import task_instance_mutation_hook
diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py
b/airflow-core/src/airflow/serialization/serialized_objects.py
index 8dd1ca10ec7..b22bc6523ae 100644
--- a/airflow-core/src/airflow/serialization/serialized_objects.py
+++ b/airflow-core/src/airflow/serialization/serialized_objects.py
@@ -51,7 +51,7 @@ from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.models.xcom import XComModel
from airflow.models.xcom_arg import SchedulerXComArg, deserialize_xcom_arg
from airflow.providers_manager import ProvidersManager
-from airflow.sdk.bases.baseoperator import BaseOperator as TaskSDKBaseOperator
+from airflow.sdk.bases.operator import BaseOperator as TaskSDKBaseOperator
from airflow.sdk.definitions.asset import (
Asset,
AssetAlias,
diff --git a/airflow-core/src/airflow/utils/task_group.py
b/airflow-core/src/airflow/utils/task_group.py
index c264274e716..27d3af1ef1c 100644
--- a/airflow-core/src/airflow/utils/task_group.py
+++ b/airflow-core/src/airflow/utils/task_group.py
@@ -32,7 +32,7 @@ MappedTaskGroup: TypeAlias =
airflow.sdk.definitions.taskgroup.MappedTaskGroup
def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False):
"""Create a nested dict representation of this TaskGroup and its children
used to construct the Graph."""
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions._internal.abstractoperator import
AbstractOperator
from airflow.sdk.definitions.mappedoperator import MappedOperator
diff --git a/dev/mypy/plugin/outputs.py b/dev/mypy/plugin/outputs.py
index 485a50cca0e..c51488e4f04 100644
--- a/dev/mypy/plugin/outputs.py
+++ b/dev/mypy/plugin/outputs.py
@@ -25,7 +25,7 @@ from mypy.types import AnyType, Type, TypeOfAny
OUTPUT_PROPERTIES = {
"airflow.models.baseoperator.BaseOperator.output",
"airflow.models.mappedoperator.MappedOperator.output",
- "airflow.sdk.bases.baseoperator.BaseOperator.output",
+ "airflow.sdk.bases.operator.BaseOperator.output",
}
TASK_CALL_FUNCTIONS = {
diff --git a/devel-common/src/tests_common/pytest_plugin.py
b/devel-common/src/tests_common/pytest_plugin.py
index d645a100d04..baed8ca2ca8 100644
--- a/devel-common/src/tests_common/pytest_plugin.py
+++ b/devel-common/src/tests_common/pytest_plugin.py
@@ -46,7 +46,7 @@ if TYPE_CHECKING:
from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk.api.datamodels._generated import IntermediateTIState,
TerminalTIState
- from airflow.sdk.bases.baseoperator import BaseOperator as
TaskSDKBaseOperator
+ from airflow.sdk.bases.operator import BaseOperator as TaskSDKBaseOperator
from airflow.sdk.execution_time.comms import StartupDetails, ToSupervisor
from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
from airflow.timetables.base import DataInterval
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py
b/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py
index 794681b4135..03183e654de 100644
---
a/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py
+++
b/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py
@@ -38,7 +38,7 @@ DISABLE_OL_PARAM = Param(False, const=False)
T = TypeVar("T", bound="DAG | Operator")
if TYPE_CHECKING:
- from airflow.sdk.bases.baseoperator import BaseOperator as SdkBaseOperator
+ from airflow.sdk.bases.operator import BaseOperator as SdkBaseOperator
log = logging.getLogger(__name__)
diff --git
a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
index 14255f28127..dde336c6a62 100644
--- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
+++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
@@ -76,7 +76,7 @@ if AIRFLOW_V_2_10_PLUS:
if AIRFLOW_V_3_0_PLUS:
from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance
as SDKTaskInstance
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.execution_time import task_runner
from airflow.sdk.execution_time.comms import StartupDetails
from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance,
parse
diff --git a/scripts/ci/pre_commit/check_base_operator_partial_arguments.py
b/scripts/ci/pre_commit/check_base_operator_partial_arguments.py
index 83a901bae99..d788e79d317 100755
--- a/scripts/ci/pre_commit/check_base_operator_partial_arguments.py
+++ b/scripts/ci/pre_commit/check_base_operator_partial_arguments.py
@@ -30,7 +30,7 @@ sys.path.insert(
from common_precommit_utils import AIRFLOW_CORE_SOURCES_PATH,
AIRFLOW_TASK_SDK_SOURCES_PATH, console
BASEOPERATOR_PY = AIRFLOW_CORE_SOURCES_PATH / "airflow" / "models" /
"baseoperator.py"
-SDK_BASEOPERATOR_PY = AIRFLOW_TASK_SDK_SOURCES_PATH / "airflow" / "sdk" /
"bases" / "baseoperator.py"
+SDK_BASEOPERATOR_PY = AIRFLOW_TASK_SDK_SOURCES_PATH / "airflow" / "sdk" /
"bases" / "operator.py"
SDK_MAPPEDOPERATOR_PY = (
AIRFLOW_TASK_SDK_SOURCES_PATH / "airflow" / "sdk" / "definitions" /
"mappedoperator.py"
)
diff --git a/task-sdk/src/airflow/sdk/__init__.py
b/task-sdk/src/airflow/sdk/__init__.py
index a7c768531b0..075b4fbef16 100644
--- a/task-sdk/src/airflow/sdk/__init__.py
+++ b/task-sdk/src/airflow/sdk/__init__.py
@@ -52,8 +52,8 @@ __all__ = [
__version__ = "1.0.0.alpha1"
if TYPE_CHECKING:
- from airflow.sdk.bases.baseoperator import BaseOperator, chain,
chain_linear, cross_downstream
from airflow.sdk.bases.notifier import BaseNotifier
+ from airflow.sdk.bases.operator import BaseOperator, chain, chain_linear,
cross_downstream
from airflow.sdk.bases.operatorlink import BaseOperatorLink
from airflow.sdk.bases.sensor import BaseSensorOperator
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAll,
AssetAny, AssetWatcher
@@ -76,7 +76,7 @@ __lazy_imports: dict[str, str] = {
"AssetAny": ".definitions.asset",
"AssetWatcher": ".definitions.asset",
"BaseNotifier": ".definitions.notifier",
- "BaseOperator": ".bases.baseoperator",
+ "BaseOperator": ".bases.operator",
"BaseOperatorLink": ".bases.operatorlink",
"BaseSensorOperator": ".bases.sensor",
"Connection": ".definitions.connection",
@@ -90,9 +90,9 @@ __lazy_imports: dict[str, str] = {
"Variable": ".definitions.variable",
"XComArg": ".definitions.xcom_arg",
"asset": ".definitions.asset.decorators",
- "chain": ".bases.baseoperator",
- "chain_linear": ".bases.baseoperator",
- "cross_downstream": ".bases.baseoperator",
+ "chain": ".bases.operator",
+ "chain_linear": ".bases.operator",
+ "cross_downstream": ".bases.operator",
"dag": ".definitions.dag",
"get_current_context": ".definitions.context",
"get_parsing_context": ".definitions.context",
diff --git a/task-sdk/src/airflow/sdk/bases/baseoperator.py
b/task-sdk/src/airflow/sdk/bases/operator.py
similarity index 100%
rename from task-sdk/src/airflow/sdk/bases/baseoperator.py
rename to task-sdk/src/airflow/sdk/bases/operator.py
diff --git a/task-sdk/src/airflow/sdk/bases/sensor.py
b/task-sdk/src/airflow/sdk/bases/sensor.py
index e896fb9f538..4347e684a42 100644
--- a/task-sdk/src/airflow/sdk/bases/sensor.py
+++ b/task-sdk/src/airflow/sdk/bases/sensor.py
@@ -37,7 +37,7 @@ from airflow.exceptions import (
TaskDeferralTimeout,
)
from airflow.executors.executor_loader import ExecutorLoader
-from airflow.sdk.bases.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
from airflow.utils import timezone
if TYPE_CHECKING:
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
index b5b68d56330..8b831d551e8 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
@@ -41,7 +41,7 @@ from airflow.utils.weight_rule import WeightRule
if TYPE_CHECKING:
import jinja2
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.bases.operatorlink import BaseOperatorLink
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.dag import DAG
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py
b/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py
index 71b88dc4000..7d0a02f6c8f 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py
@@ -22,7 +22,7 @@ from collections.abc import Iterable, Sequence
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions._internal.abstractoperator import Operator
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.edges import EdgeModifier
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/node.py
b/task-sdk/src/airflow/sdk/definitions/_internal/node.py
index 7fab2f1919b..05968290d7b 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/node.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/node.py
@@ -122,7 +122,7 @@ class DAGNode(DependencyMixin, metaclass=ABCMeta):
edge_modifier: EdgeModifier | None = None,
) -> None:
"""Set relatives for the task or task list."""
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.mappedoperator import MappedOperator
if not isinstance(task_or_task_list, Sequence):
diff --git a/task-sdk/src/airflow/sdk/definitions/context.py
b/task-sdk/src/airflow/sdk/definitions/context.py
index 03fbb6a7763..c54a5edd6a8 100644
--- a/task-sdk/src/airflow/sdk/definitions/context.py
+++ b/task-sdk/src/airflow/sdk/definitions/context.py
@@ -25,7 +25,7 @@ if TYPE_CHECKING:
from datetime import datetime
from airflow.models.operator import Operator
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.dag import DAG
from airflow.sdk.execution_time.context import InletEventsAccessors
from airflow.sdk.types import (
diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py
b/task-sdk/src/airflow/sdk/definitions/dag.py
index dd2f4cef71d..ef2f0b0d0bd 100644
--- a/task-sdk/src/airflow/sdk/definitions/dag.py
+++ b/task-sdk/src/airflow/sdk/definitions/dag.py
@@ -51,7 +51,7 @@ from airflow.exceptions import (
ParamValidationError,
TaskNotFound,
)
-from airflow.sdk.bases.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
from airflow.sdk.definitions._internal.types import NOTSET
from airflow.sdk.definitions.asset import AssetAll, BaseAsset
diff --git a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
index 7416967d42c..9fed450ca75 100644
--- a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
@@ -69,7 +69,7 @@ if TYPE_CHECKING:
OperatorExpandKwargsArgument,
)
from airflow.models.xcom_arg import XComArg
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.bases.operatorlink import BaseOperatorLink
from airflow.sdk.definitions.dag import DAG
from airflow.sdk.definitions.param import ParamsDict
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index c03c50cc4cd..65ae3f4c2c4 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -41,7 +41,7 @@ from airflow.utils.trigger_rule import TriggerRule
if TYPE_CHECKING:
from airflow.models.expandinput import ExpandInput
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions._internal.abstractoperator import
AbstractOperator
from airflow.sdk.definitions._internal.mixins import DependencyMixin
from airflow.sdk.definitions.dag import DAG
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index c5c124cf3c8..d5411b9dd8e 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -34,7 +34,7 @@ from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.xcom import XCOM_RETURN_KEY
if TYPE_CHECKING:
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.edges import EdgeModifier
from airflow.sdk.types import Operator
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index 5ce223962a2..843ec37a24a 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -43,7 +43,7 @@ if TYPE_CHECKING:
from uuid import UUID
from airflow.sdk import Variable
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.connection import Connection
from airflow.sdk.definitions.context import Context
from airflow.sdk.execution_time.comms import (
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 4d5e9dcd753..747f135389c 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -46,7 +46,7 @@ from airflow.sdk.api.datamodels._generated import (
TerminalTIState,
TIRunContext,
)
-from airflow.sdk.bases.baseoperator import BaseOperator, ExecutorSafeguard
+from airflow.sdk.bases.operator import BaseOperator, ExecutorSafeguard
from airflow.sdk.definitions._internal.dag_parsing_context import
_airflow_parsing_context_manager
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetNameRef,
AssetUniqueKey, AssetUriRef
diff --git a/task-sdk/src/airflow/sdk/types.py
b/task-sdk/src/airflow/sdk/types.py
index e7e2036aa5c..cd31c4e71c8 100644
--- a/task-sdk/src/airflow/sdk/types.py
+++ b/task-sdk/src/airflow/sdk/types.py
@@ -29,7 +29,7 @@ if TYPE_CHECKING:
from collections.abc import Iterator
from datetime import datetime
- from airflow.sdk.bases.baseoperator import BaseOperator
+ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.asset import Asset, AssetAlias,
AssetAliasEvent, AssetRef, BaseAssetUniqueKey
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.mappedoperator import MappedOperator
diff --git a/task-sdk/tests/task_sdk/bases/test_baseoperator.py
b/task-sdk/tests/task_sdk/bases/test_baseoperator.py
index 969030a22ca..b2c1440aeb8 100644
--- a/task-sdk/tests/task_sdk/bases/test_baseoperator.py
+++ b/task-sdk/tests/task_sdk/bases/test_baseoperator.py
@@ -29,7 +29,7 @@ import pytest
import structlog
from airflow.decorators import task as task_decorator
-from airflow.sdk.bases.baseoperator import (
+from airflow.sdk.bases.operator import (
BaseOperator,
BaseOperatorMeta,
ExecutorSafeguard,
diff --git a/task-sdk/tests/task_sdk/dags/super_basic.py
b/task-sdk/tests/task_sdk/dags/super_basic.py
index 2cccb9ab4c6..cfa2e28475c 100644
--- a/task-sdk/tests/task_sdk/dags/super_basic.py
+++ b/task-sdk/tests/task_sdk/dags/super_basic.py
@@ -17,7 +17,7 @@
from __future__ import annotations
-from airflow.sdk.bases.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.dag import dag
diff --git a/task-sdk/tests/task_sdk/dags/super_basic_run.py
b/task-sdk/tests/task_sdk/dags/super_basic_run.py
index e178e7acc89..7f385ddf393 100644
--- a/task-sdk/tests/task_sdk/dags/super_basic_run.py
+++ b/task-sdk/tests/task_sdk/dags/super_basic_run.py
@@ -17,7 +17,7 @@
from __future__ import annotations
-from airflow.sdk.bases.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.dag import dag
diff --git a/task-sdk/tests/task_sdk/definitions/test_dag.py
b/task-sdk/tests/task_sdk/definitions/test_dag.py
index 07a99221d49..6c00f69653b 100644
--- a/task-sdk/tests/task_sdk/definitions/test_dag.py
+++ b/task-sdk/tests/task_sdk/definitions/test_dag.py
@@ -23,7 +23,7 @@ from typing import Any
import pytest
from airflow.exceptions import DuplicateTaskIdFound
-from airflow.sdk.bases.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.dag import DAG, dag as dag_decorator
from airflow.sdk.definitions.param import DagParam, Param, ParamsDict
diff --git a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
index 3c780551c6a..925f087a282 100644
--- a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
+++ b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
@@ -25,7 +25,7 @@ import pendulum
import pytest
from airflow.sdk.api.datamodels._generated import TerminalTIState
-from airflow.sdk.bases.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.dag import DAG
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.xcom_arg import XComArg
diff --git a/task-sdk/tests/task_sdk/definitions/test_mixins.py
b/task-sdk/tests/task_sdk/definitions/test_mixins.py
index b7ffa974758..b5e5e4bfe26 100644
--- a/task-sdk/tests/task_sdk/definitions/test_mixins.py
+++ b/task-sdk/tests/task_sdk/definitions/test_mixins.py
@@ -22,7 +22,7 @@ import itertools
import pytest
from airflow.decorators import setup, task, teardown
-from airflow.sdk.bases.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.dag import DAG
