Skip to content

Commit 8d18c69

Browse files
authored
Add type casting for mypy to understand expand_input types (#53191)
* Add type casting for mypy to understand upstream_mapped_group._expand_input * Add type casting for mypy to understand upstream_mapped_group._expand_input
1 parent 5aec2d5 commit 8d18c69

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from collections import defaultdict
2424
from collections.abc import Iterator
2525
from datetime import datetime
26-
from typing import TYPE_CHECKING, Annotated, Any
26+
from typing import TYPE_CHECKING, Annotated, Any, cast
2727
from uuid import UUID
2828

2929
import attrs
@@ -73,6 +73,7 @@
7373
if TYPE_CHECKING:
7474
from sqlalchemy.sql.dml import Update
7575

76+
from airflow.models.expandinput import SchedulerExpandInput
7677
from airflow.sdk.types import Operator
7778

7879

@@ -308,9 +309,9 @@ def _get_upstream_map_indexes(
308309
mapped_ti_count = upstream_mapped_group.get_parse_time_mapped_ti_count()
309310
except NotFullyPopulated:
310311
# for cases that needs to resolve xcom to get the correct count
311-
mapped_ti_count = upstream_mapped_group._expand_input.get_total_map_length(
312-
run_id, session=session
313-
)
312+
mapped_ti_count = cast(
313+
"SchedulerExpandInput", upstream_mapped_group._expand_input
314+
).get_total_map_length(run_id, session=session)
314315
map_indexes = list(range(mapped_ti_count)) if mapped_ti_count is not None else None
315316

316317
yield upstream_task.task_id, map_indexes

0 commit comments

Comments
 (0)