Skip to content
Merged
Changes from all commits
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 @@ -73,6 +73,7 @@
if TYPE_CHECKING:
from sqlalchemy.sql.dml import Update

from airflow.models.expandinput import SchedulerExpandInput
from airflow.sdk.types import Operator


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)
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