Skip to content

Commit

Permalink
Merge pull request #2032 from PrefectHQ/issue-1821-add-better-warning
Browse files Browse the repository at this point in the history
Add early check for mapping argument type errors
  • Loading branch information
cicdw committed Feb 15, 2020
2 parents 251aa4f + 1e68171 commit ee383b8
Show file tree
Hide file tree
Showing 4 changed files with 22 additions and 1 deletion.
2 changes: 1 addition & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/

### Enhancements

- None
- Better exception for unsubscriptable mapping arguments - [#1821](https://github.com/PrefectHQ/prefect/issues/1821)

### Task Library

Expand Down
8 changes: 8 additions & 0 deletions src/prefect/core/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import prefect.engine.signals
import prefect.triggers
from prefect.utilities import logging
from prefect.utilities.tasks import unmapped
from prefect.utilities.notifications import callback_factory

if TYPE_CHECKING:
Expand Down Expand Up @@ -475,6 +476,13 @@ def map(
Returns:
- Task: a new Task instance
"""
for arg in args:
if not hasattr(arg, "__getitem__") and not isinstance(arg, unmapped):
raise TypeError(
"Cannot map over unsubscriptable object of type {t}: {preview}...".format(
t=type(arg), preview=repr(arg)[:10]
)
)
new = self.copy(**(task_args or {}))
return new.bind(
*args, mapped=True, upstream_tasks=upstream_tasks, flow=flow, **kwargs
Expand Down
7 changes: 7 additions & 0 deletions src/prefect/engine/task_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -717,6 +717,13 @@ def run_mapped_task(
# Therefore, we only try to get a result if EITHER this task's
# state is not already mapped OR the upstream result is not None.
if not state.is_mapped() or upstream_state._result != NoResult:
if not hasattr(upstream_state.result, "__getitem__"):
raise TypeError(
"Cannot map over unsubscriptable object of type {t}: {preview}...".format(
t=type(upstream_state.result),
preview=repr(upstream_state.result)[:10],
)
)
upstream_result = Result(
upstream_state.result[i],
result_handler=upstream_state._result.result_handler, # type: ignore
Expand Down
6 changes: 6 additions & 0 deletions tests/core/test_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -581,3 +581,9 @@ def test_tags_are_appended_to_when_updating_with_task_args(self):
t2 = t(1, 2, task_args={"name": "test-tags", "tags": ["new-tag"]})

assert t2.tags == {"math", "test", "new-tag"}

def test_task_check_mapped_args_are_subscriptable_in_advance(self):
t = Task()
with pytest.raises(TypeError):
with Flow(name="test") as f:
res = t.map({1, 2, 3, 4})

0 comments on commit ee383b8

Please sign in to comment.