From 3d6e21a920de20a4343dc818325cd33f7e6f7db8 Mon Sep 17 00:00:00 2001 From: Ishaan Samantray Date: Mon, 1 Jun 2026 15:29:18 -0400 Subject: [PATCH] fix(iterator): raise TypeTransformerFailedError with a message to prevent IndexError in callers IteratorTransformer.to_python_value previously raised TypeTransformerFailedError() with no arguments. Callers in promise.py, base_task.py, and base_connector.py all access exc.args[0] to enrich the error message; with an empty args tuple this caused an IndexError that swallowed the original type conversion failure entirely. Replace the bare except/re-raise with an explicit None-check on lv.collection and raise TypeTransformerFailedError with a descriptive message, ensuring exc.args[0] is always valid. --- flytekit/types/iterator/iterator.py | 9 +++++---- .../flytekit/unit/types/iterator/test_iterator.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 4 deletions(-) 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