diff --git a/flytekit/types/iterator/iterator.py b/flytekit/types/iterator/iterator.py index cdd532f9b7..1a46f59953 100644 --- a/flytekit/types/iterator/iterator.py +++ b/flytekit/types/iterator/iterator.py @@ -57,10 +57,11 @@ def to_literal( return Literal(collection=LiteralCollection(literals=lit_list)) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Type[T]) -> FlyteIterator: - try: - lits = lv.collection.literals - except AttributeError: - raise TypeTransformerFailedError() + if lv.collection is None: + raise TypeTransformerFailedError( + f"Expected a collection literal for {expected_python_type}, but got a non-collection value." + ) + lits = lv.collection.literals return FlyteIterator(ctx, lv, expected_python_type, len(lits)) diff --git a/tests/flytekit/unit/types/iterator/test_iterator.py b/tests/flytekit/unit/types/iterator/test_iterator.py index 53957e639e..e14787889a 100644 --- a/tests/flytekit/unit/types/iterator/test_iterator.py +++ b/tests/flytekit/unit/types/iterator/test_iterator.py @@ -1,6 +1,12 @@ import typing +import pytest + from flytekit import task, workflow +from flytekit.core.context_manager import FlyteContextManager +from flytekit.core.type_engine import TypeTransformerFailedError +from flytekit.models.literals import Literal, Scalar, Primitive +from flytekit.types.iterator.iterator import IteratorTransformer @task @@ -21,3 +27,12 @@ def wf(a: int) -> typing.List[int]: def test_iterator(): assert wf(a=4) == [0, 1, 2, 3] + + +def test_to_python_value_non_collection_raises_with_message(): + ctx = FlyteContextManager.current_context() + lit = Literal(scalar=Scalar(primitive=Primitive(integer=42))) + trans = IteratorTransformer() + with pytest.raises(TypeTransformerFailedError) as exc_info: + trans.to_python_value(ctx, lit, typing.Iterator[int]) + assert exc_info.value.args[0] # args[0] must be set — not empty tuple