Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
from collections import defaultdict
from collections.abc import Iterator
from datetime import datetime
from typing import TYPE_CHECKING, Annotated, Any
from typing import TYPE_CHECKING, Annotated, Any, cast
from uuid import UUID

import attrs
Expand Down Expand Up @@ -60,6 +60,7 @@
from airflow.models.asset import AssetActive
from airflow.models.dagbag import DagBag
from airflow.models.dagrun import DagRun as DR
from airflow.models.expandinput import SchedulerExpandInput
Comment thread
gopidesupavan marked this conversation as resolved.
Outdated
from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
from airflow.models.taskreschedule import TaskReschedule
from airflow.models.trigger import Trigger
Expand Down Expand Up @@ -308,9 +309,9 @@ def _get_upstream_map_indexes(
mapped_ti_count = upstream_mapped_group.get_parse_time_mapped_ti_count()
except NotFullyPopulated:
# for cases that needs to resolve xcom to get the correct count
mapped_ti_count = upstream_mapped_group._expand_input.get_total_map_length(
run_id, session=session
)
mapped_ti_count = cast(
"SchedulerExpandInput", upstream_mapped_group._expand_input
).get_total_map_length(run_id, session=session)
Comment thread
gopidesupavan marked this conversation as resolved.
map_indexes = list(range(mapped_ti_count)) if mapped_ti_count is not None else None

yield upstream_task.task_id, map_indexes
Expand Down