Re: [PR] fix(task_instances): handle upstream_mapped_index when xcom access is needed [airflow]

["uranusjr (via GitHub)" <gi...@apache.org>]

uranusjr commented on code in PR #50641:
URL: https://github.com/apache/airflow/pull/50641#discussion_r2101429445


##########
airflow-core/src/airflow/models/expandinput.py:
##########
@@ -130,11 +133,13 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int:
         return length
 
 
-_EXPAND_INPUT_TYPES = {
+_EXPAND_INPUT_TYPES: dict[str, type] = {
     "dict-of-lists": SchedulerDictOfListsExpandInput,
     "list-of-dicts": SchedulerListOfDictsExpandInput,
 }
 
+SchedulerExpandInput = Union[SchedulerDictOfListsExpandInput, SchedulerListOfDictsExpandInput]
+
 
-def create_expand_input(kind: str, value: Any) -> ExpandInput:
+def create_expand_input(kind: str, value: Any) -> SchedulerExpandInput:
     return _EXPAND_INPUT_TYPES[kind](value)

Review Comment:
   How about just create an ExpandInput base class for both kinds to subclass, and use it here instead of the union?



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org