Repository: apache/airflow

Issue:
expand_kwargs.map(func) gives unhelpful error message if func returns list
### Apache Airflow version

main (development)

### What happened

Here's a DAG:

```python3

with DAG(
    dag_id="expand_list",
    doc_md="try to get kwargs from a list",
    schedule_interval=None,
    start_date=datetime(2001, 1, 1),
) as expand_list:

    @expand_list.task
    def do_this():
        return [
            ("echo hello $USER", "USER", "foo"),
            ("echo hello $USER", "USER", "bar"),
        ]

    def mapper(tuple):
        if tuple[2] == "bar":
            return [1, 2, 3]
        else:
            return {"bash_command": tuple[0], "env": {tuple[1]: tuple[2]}}

    BashOperator.partial(task_id="one_cmd").expand_kwargs(do_this().map(mapper))
```

The `foo` task instance succeeds  as expected, and the `bar` task fails as expected.  But the error message that it gives isn't particularly helpful to a user who doesn't know what they did wrong:

```
ERROR - Failed to execute task: resolve() takes 3 positional arguments but 4 were given.
Traceback (most recent call last):
  File "/home/matt/src/airflow/airflow/executors/debug_executor.py", line 78, in _run_task
    ti.run(job_id=ti.job_id, **params)
  File "/home/matt/src/airflow/airflow/utils/session.py", line 71, in wrapper
    return func(*args, session=session, **kwargs)
  File "/home/matt/src/airflow/airflow/models/taskinstance.py", line 1782, in run
    self._run_raw_task(
  File "/home/matt/src/airflow/airflow/utils/session.py", line 68, in wrapper
    return func(*args, **kwargs)
  File "/home/matt/src/airflow/airflow/models/taskinstance.py", line 1445, in _run_raw_task
    self._execute_task_with_callbacks(context, test_mode)
  File "/home/matt/src/airflow/airflow/models/taskinstance.py", line 1580, in _execute_task_with_callbacks
    task_orig = self.render_templates(context=context)
  File "/home/matt/src/airflow/airflow/models/taskinstance.py", line 2202, in render_templates
    rendered_task = self.task.render_template_fields(context)
  File "/home/matt/src/airflow/airflow/models/mappedoperator.py", line 751, in render_template_fields
    unmapped_task = self.unmap(mapped_kwargs)
  File "/home/matt/src/airflow/airflow/models/mappedoperator.py", line 591, in unmap
    kwargs = self._expand_mapped_kwargs(resolve)
  File "/home/matt/src/airflow/airflow/models/mappedoperator.py", line 546, in _expand_mapped_kwargs
    return expand_input.resolve(*resolve)
TypeError: resolve() takes 3 positional arguments but 4 were given
```

### What you think should happen instead

Whatever checks the return value for mappability should do more to point the user to their error.  Perhaps something like:

> UnmappableDataError: Expected a dict with keys that BashOperator accepts, got `[1, 2, 3]` instead

### How to reproduce

Run the dag above

### Operating System

Linux 5.10.101 #1-NixOS SMP Wed Feb 16 11:54:31 UTC 2022 x86_64 GNU/Linux

### Versions of Apache Airflow Providers

n/a

### Deployment

Virtualenv installation

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:

Resolving patch:
diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -455,7 +455,7 @@ def _expand_mapped_kwargs(self, resolve: Optional[Tuple[Context, Session]]) -> D
         assert self.expand_input is EXPAND_INPUT_EMPTY
         return {"op_kwargs": super()._expand_mapped_kwargs(resolve)}
 
-    def _get_unmap_kwargs(self, mapped_kwargs: Dict[str, Any], *, strict: bool) -> Dict[str, Any]:
+    def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -> Dict[str, Any]:
         if strict:
             prevent_duplicates(
                 self.partial_kwargs["op_kwargs"],
diff --git a/airflow/models/expandinput.py b/airflow/models/expandinput.py
--- a/airflow/models/expandinput.py
+++ b/airflow/models/expandinput.py
@@ -22,7 +22,7 @@
 import collections.abc
 import functools
 import operator
-from typing import TYPE_CHECKING, Any, Iterable, NamedTuple, Sequence, Sized, Union
+from typing import TYPE_CHECKING, Any, Iterable, Mapping, NamedTuple, Sequence, Sized, Union
 
 from sqlalchemy import func
 from sqlalchemy.orm import Session
@@ -195,10 +195,16 @@ def _find_index_for_this_field(index: int) -> int:
                 return k, v
         raise IndexError(f"index {map_index} is over mapped length")
 
-    def resolve(self, context: Context, session: Session) -> dict[str, Any]:
+    def resolve(self, context: Context, session: Session) -> Mapping[str, Any]:
         return {k: self._expand_mapped_field(k, v, context, session=session) for k, v in self.value.items()}
 
 
+def _describe_type(value: Any) -> str:
+    if value is None:
+        return "None"
+    return type(value).__name__
+
+
 class ListOfDictsExpandInput(NamedTuple):
     """Storage type of a mapped operator's mapped kwargs.
 
@@ -245,12 +251,23 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int:
             raise NotFullyPopulated({"expand_kwargs() argument"})
         return value
 
-    def resolve(self, context: Context, session: Session) -> dict[str, Any]:
+    def resolve(self, context: Context, session: Session) -> Mapping[str, Any]:
         map_index = context["ti"].map_index
         if map_index < 0:
             raise RuntimeError("can't resolve task-mapping argument without expanding")
-        # Validation should be done when the upstream returns.
-        return self.value.resolve(context, session)[map_index]
+        mappings = self.value.resolve(context, session)
+        if not isinstance(mappings, collections.abc.Sequence):
+            raise ValueError(f"expand_kwargs() expects a list[dict], not {_describe_type(mappings)}")
+        mapping = mappings[map_index]
+        if not isinstance(mapping, collections.abc.Mapping):
+            raise ValueError(f"expand_kwargs() expects a list[dict], not list[{_describe_type(mapping)}]")
+        for key in mapping:
+            if not isinstance(key, str):
+                raise ValueError(
+                    f"expand_kwargs() input dict keys must all be str, "
+                    f"but {key!r} is of type {_describe_type(key)}"
+                )
+        return mapping
 
 
 EXPAND_INPUT_EMPTY = DictOfListsExpandInput({})  # Sentinel value.
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -30,6 +30,7 @@
     Iterable,
     Iterator,
     List,
+    Mapping,
     Optional,
     Sequence,
     Set,
@@ -123,7 +124,7 @@ def validate_mapping_kwargs(op: Type["BaseOperator"], func: ValidationSource, va
     raise TypeError(f"{op.__name__}.{func}() got {error}")
 
 
-def prevent_duplicates(kwargs1: Dict[str, Any], kwargs2: Dict[str, Any], *, fail_reason: str) -> None:
+def prevent_duplicates(kwargs1: Dict[str, Any], kwargs2: Mapping[str, Any], *, fail_reason: str) -> None:
     duplicated_keys = set(kwargs1).intersection(kwargs2)
     if not duplicated_keys:
         return
@@ -528,7 +529,7 @@ def serialize_for_task_group(self) -> Tuple[DagAttributeTypes, Any]:
         """Implementing DAGNode."""
         return DagAttributeTypes.OP, self.task_id
 
-    def _expand_mapped_kwargs(self, resolve: Optional[Tuple[Context, Session]]) -> Dict[str, Any]:
+    def _expand_mapped_kwargs(self, resolve: Optional[Tuple[Context, Session]]) -> Mapping[str, Any]:
         """Get the kwargs to create the unmapped operator.
 
         If *resolve* is not *None*, it must be a two-tuple to provide context to
@@ -546,7 +547,7 @@ def _expand_mapped_kwargs(self, resolve: Optional[Tuple[Context, Session]]) -> D
             return expand_input.resolve(*resolve)
         return expand_input.get_unresolved_kwargs()
 
-    def _get_unmap_kwargs(self, mapped_kwargs: Dict[str, Any], *, strict: bool) -> Dict[str, Any]:
+    def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -> Dict[str, Any]:
         """Get init kwargs to unmap the underlying operator class.
 
         :param mapped_kwargs: The dict returned by ``_expand_mapped_kwargs``.
@@ -569,7 +570,7 @@ def _get_unmap_kwargs(self, mapped_kwargs: Dict[str, Any], *, strict: bool) -> D
             **mapped_kwargs,
         }
 
-    def unmap(self, resolve: Union[None, Dict[str, Any], Tuple[Context, Session]]) -> "BaseOperator":
+    def unmap(self, resolve: Union[None, Mapping[str, Any], Tuple[Context, Session]]) -> "BaseOperator":
         """Get the "normal" Operator after applying the current mapping.
 
         If ``operator_class`` is not a class (i.e. this DAG has been


Test patch:
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -34,6 +34,7 @@
 from freezegun import freeze_time
 
 from airflow import models, settings
+from airflow.decorators import task
 from airflow.example_dags.plugins.workday import AfterWorkdayTimetable
 from airflow.exceptions import (
     AirflowException,
@@ -2640,6 +2641,41 @@ def push():
         assert ti.state == TaskInstanceState.FAILED
         assert str(ctx.value) == error_message
 
+    @pytest.mark.parametrize(
+        "create_upstream",
+        [
+            # The task returns an invalid expand_kwargs() input (a list[int] instead of list[dict]).
+            pytest.param(lambda: task(task_id="push")(lambda: [0])(), id="normal"),
+            # This task returns a list[dict] (correct), but we use map() to transform it to list[int] (wrong).
+            pytest.param(lambda: task(task_id="push")(lambda: [{"v": ""}])().map(lambda _: 0), id="mapped"),
+        ],
+    )
+    def test_expand_kwargs_error_if_received_invalid(self, dag_maker, session, create_upstream):
+        with dag_maker(dag_id="test_expand_kwargs_error_if_received_invalid", session=session):
+            push_task = create_upstream()
+
+            @task()
+            def pull(v):
+                print(v)
+
+            pull.expand_kwargs(push_task)
+
+        dr = dag_maker.create_dagrun()
+
+        # Run "push".
+        decision = dr.task_instance_scheduling_decisions(session=session)
+        assert decision.schedulable_tis
+        for ti in decision.schedulable_tis:
+            ti.run()
+
+        # Run "pull".
+        decision = dr.task_instance_scheduling_decisions(session=session)
+        assert decision.schedulable_tis
+        for ti in decision.schedulable_tis:
+            with pytest.raises(ValueError) as ctx:
+                ti.run()
+            assert str(ctx.value) == "expand_kwargs() expects a list[dict], not list[int]"
+
     @pytest.mark.parametrize(
         "downstream, error_message",
         [
diff --git a/tests/models/test_xcom_arg_map.py b/tests/models/test_xcom_arg_map.py
--- a/tests/models/test_xcom_arg_map.py
+++ b/tests/models/test_xcom_arg_map.py
@@ -128,9 +128,9 @@ def c_to_none(v):
     tis[("pull", 1)].run()
 
     # But the third one fails because the map() result cannot be used as kwargs.
-    with pytest.raises(TypeError) as ctx:
+    with pytest.raises(ValueError) as ctx:
         tis[("pull", 2)].run()
-    assert str(ctx.value) == "'NoneType' object is not iterable"
+    assert str(ctx.value) == "expand_kwargs() expects a list[dict], not list[None]"
 
     assert [tis[("pull", i)].state for i in range(3)] == [
         TaskInstanceState.SUCCESS,