uranusjr commented on code in PR #52234:
URL: https://github.com/apache/airflow/pull/52234#discussion_r2196869199
##########
airflow-core/src/airflow/models/mappedoperator.py:
##########
@@ -173,3 +146,87 @@ 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)}")
+
+
+# Still accept TaskSDKBaseOperator because some tests don't go through
serialization.
+# TODO (GH-52141): Rewrite tests so we can drop SDK references at some point.
+@get_mapped_ti_count.register(SerializedBaseOperator)
+@get_mapped_ti_count.register(TaskSDKBaseOperator)
+def _(task: SerializedBaseOperator | TaskSDKBaseOperator, run_id: str, *,
session: Session) -> int:
Review Comment:
Turns out `singledispatch` does not support unions so we’ll need to keep
this verbose syntax a bit longer.
--
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]
