Remove MappedOperator inheritance

Now MappedOperator in core is entirely detached from the SDK.

I decided to remove the start_from_trigger code entirely (the code path
is already entirely unreachable in Airflow 3) since it causes too many
Mypy errors. We can always get the implementation back from Git history
when we re-implement the feature with the new task runner.

Author: uranusjr

airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py

@@ -23,8 +23,8 @@
 import structlog
 
 from airflow.api_fastapi.common.parameters import state_priority
+from airflow.models.mappedoperator import MappedOperator
 from airflow.models.taskmap import TaskMap
-from airflow.sdk.definitions.mappedoperator import MappedOperator
 from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup, get_task_group_children_getter
 from airflow.serialization.serialized_objects import SerializedBaseOperator
 

airflow-core/src/airflow/models/dagrun.py

@@ -102,13 +102,11 @@
     from airflow.serialization.serialized_objects import SerializedBaseOperator
     from airflow.utils.types import ArgNotSet
 
-    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
-
     CreatedTasks = TypeVar("CreatedTasks", Iterator["dict[str, Any]"], Iterator[TI])
-
-    AttributeValueType = (
+    AttributeValueType: TypeAlias = (
         str | bool | int | float | Sequence[str] | Sequence[bool] | Sequence[int] | Sequence[float]
     )
+    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
 
 RUN_ID_REGEX = r"^(?:manual|scheduled|asset_triggered)__(?:\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\+00:00)$"
 
@@ -1483,15 +1481,15 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
             If the ti does not need expansion, either because the task is not
             mapped, or has already been expanded, *None* is returned.
             """
+            from airflow.models.mappedoperator import is_mapped
+
             if TYPE_CHECKING:
                 assert ti.task
 
             if ti.map_index >= 0:  # Already expanded, we're good.
                 return None
 
-            from airflow.sdk.definitions.mappedoperator import MappedOperator as TaskSDKMappedOperator
-
-            if isinstance(ti.task, TaskSDKMappedOperator):
+            if is_mapped(ti.task):
                 # If we get here, it could be that we are moving from non-mapped to mapped
                 # after task instance clearing or this ti is not yet expanded. Safe to clear
                 # the db references.
@@ -1510,7 +1508,7 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
         revised_map_index_task_ids: set[str] = set()
         for schedulable in itertools.chain(schedulable_tis, additional_tis):
             if TYPE_CHECKING:
-                assert isinstance(schedulable.task, SerializedBaseOperator)
+                assert isinstance(schedulable.task, Operator)
             old_state = schedulable.state
             if not schedulable.are_dependencies_met(session=session, dep_context=dep_context):
                 old_states[schedulable.key] = old_state
@@ -1995,25 +1993,23 @@ def schedule_tis(
         empty_ti_ids: list[str] = []
         schedulable_ti_ids: list[str] = []
         for ti in schedulable_tis:
+            task = ti.task
             if TYPE_CHECKING:
-                assert isinstance(ti.task, SerializedBaseOperator)
+                assert isinstance(task, Operator)
             if (
-                ti.task.inherits_from_empty_operator
-                and not ti.task.on_execute_callback
-                and not ti.task.on_success_callback
-                and not ti.task.outlets
-                and not ti.task.inlets
+                task.inherits_from_empty_operator
+                and not task.on_execute_callback
+                and not task.on_success_callback
+                and not task.outlets
+                and not task.inlets
             ):
                 empty_ti_ids.append(ti.id)
             # check "start_trigger_args" to see whether the operator supports start execution from triggerer
             # if so, we'll then check "start_from_trigger" to see whether this feature is turned on and defer
             # this task.
             # if not, we'll add this "ti" into "schedulable_ti_ids" and later execute it to run in the worker
-            elif ti.task.start_trigger_args is not None:
-                context = ti.get_template_context()
-                start_from_trigger = ti.task.expand_start_from_trigger(context=context, session=session)
-
-                if start_from_trigger:
+            elif task.start_trigger_args is not None:
+                if task.expand_start_from_trigger(context=ti.get_template_context()):
                     ti.start_date = timezone.utcnow()
                     if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:
                         ti.try_number += 1

airflow-core/src/airflow/models/expandinput.py

@@ -19,18 +19,11 @@
 
 import functools
 import operator
-from collections.abc import Iterable, Sized
+from collections.abc import Iterable, Mapping, Sequence, Sized
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import attrs
 
-if TYPE_CHECKING:
-    from typing import TypeGuard
-
-    from sqlalchemy.orm import Session
-
-    from airflow.models.xcom_arg import SchedulerXComArg
-
 from airflow.sdk.definitions._internal.expandinput import (
     DictOfListsExpandInput,
     ListOfDictsExpandInput,
@@ -41,6 +34,18 @@
     is_mappable,
 )
 
+if TYPE_CHECKING:
+    from typing import TypeAlias, TypeGuard
+
+    from sqlalchemy.orm import Session
+
+    from airflow.models.mappedoperator import MappedOperator
+    from airflow.models.xcom_arg import SchedulerXComArg
+    from airflow.serialization.serialized_objects import SerializedBaseOperator
+
+    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
+
+
 __all__ = [
     "DictOfListsExpandInput",
     "ListOfDictsExpandInput",
@@ -111,6 +116,26 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int:
         lengths = self._get_map_lengths(run_id, session=session)
         return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1)
 
+    def iter_references(self) -> Iterable[tuple[Operator, str]]:
+        from airflow.models.referencemixin import ReferenceMixin
+
+        for x in self.value.values():
+            if isinstance(x, ReferenceMixin):
+                yield from x.iter_references()
+
+
+# To replace tedious isinstance() checks.
+def _is_parse_time_mappable(v: OperatorExpandArgument) -> TypeGuard[Mapping | Sequence]:
+    from airflow.sdk.definitions.xcom_arg import XComArg
+
+    return not isinstance(v, (MappedArgument, XComArg))
+
+
+def _describe_type(value: Any) -> str:
+    if value is None:
+        return "None"
+    return type(value).__name__
+
 
 @attrs.define
 class SchedulerListOfDictsExpandInput:
@@ -133,6 +158,16 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int:
             raise NotFullyPopulated({"expand_kwargs() argument"})
         return length
 
+    def iter_references(self) -> Iterable[tuple[Operator, str]]:
+        from airflow.models.referencemixin import ReferenceMixin
+
+        if isinstance(self.value, ReferenceMixin):
+            yield from self.value.iter_references()
+        else:
+            for x in self.value:
+                if isinstance(x, ReferenceMixin):
+                    yield from x.iter_references()
+
 
 _EXPAND_INPUT_TYPES: dict[str, type[SchedulerExpandInput]] = {
     "dict-of-lists": SchedulerDictOfListsExpandInput,