kaxil commented on code in PR #52234:
URL: https://github.com/apache/airflow/pull/52234#discussion_r2190958032
##########
airflow-core/src/airflow/models/mappedoperator.py:
##########
@@ -173,3 +151,88 @@ def get_extra_links(self, ti: TaskInstance, name: str) ->
str | None:
if not link:
return None
return link.get_link(self, ti_key=ti.key) # type: ignore[arg-type]
+
+
[email protected]
+def get_mapped_ti_count(task: DAGNode, run_id: str, *, session: Session) ->
int:
+ raise NotImplementedError(f"Not implemented for {type(task)}")
+
+
+# https://github.com/python/cpython/issues/86153
+# While we support Python 3.9 we can't rely on the type hint, we need to pass
the type explicitly to
+# register.
+@get_mapped_ti_count.register(SerializedBaseOperator)
+@get_mapped_ti_count.register(TaskSDKBaseOperator) # Some tests don't go
through task serialization...
+def _(task: SerializedBaseOperator | TaskSDKBaseOperator, run_id: str, *,
session: Session) -> int:
+ group = task.get_closest_mapped_task_group()
+ if group is None:
+ raise NotMapped()
+ return get_mapped_ti_count(group, run_id, session=session)
+
+
+@get_mapped_ti_count.register(MappedOperator)
+@get_mapped_ti_count.register(TaskSDKMappedOperator) # Some tests don't go
through task serialization...
+def _(task: MappedOperator | TaskSDKMappedOperator, run_id: str, *, session:
Session) -> int:
+ from airflow.serialization.serialized_objects import BaseSerialization,
_ExpandInputRef
+
+ exp_input = task._get_specified_expand_input()
+ if isinstance(exp_input, _ExpandInputRef):
+ exp_input = exp_input.deref(task.dag)
+ # TODO: TaskSDK This is only needed to support `dag.test()` etc until we
port it over to use the
+ # task sdk runner.
+ if not hasattr(exp_input, "get_total_map_length"):
Review Comment:
If it isn't just moving the code: fine to do it in separate PR too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
