Skip to content

Commit 8fdfbd5

Browse files
committed
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.
1 parent d42ce0f commit 8fdfbd5

2 files changed

Lines changed: 20 additions & 4 deletions

File tree

flytekit/types/iterator/iterator.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -57,10 +57,11 @@ def to_literal(
5757
return Literal(collection=LiteralCollection(literals=lit_list))
5858

5959
def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Type[T]) -> FlyteIterator:
60-
try:
61-
lits = lv.collection.literals
62-
except AttributeError:
63-
raise TypeTransformerFailedError()
60+
if lv.collection is None:
61+
raise TypeTransformerFailedError(
62+
f"Expected a collection literal for {expected_python_type}, but got a non-collection value."
63+
)
64+
lits = lv.collection.literals
6465
return FlyteIterator(ctx, lv, expected_python_type, len(lits))
6566

6667

tests/flytekit/unit/types/iterator/test_iterator.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
11
import typing
22

3+
import pytest
4+
35
from flytekit import task, workflow
6+
from flytekit.core.context_manager import FlyteContextManager
7+
from flytekit.core.type_engine import TypeTransformerFailedError
8+
from flytekit.models.literals import Literal, Scalar, Primitive
9+
from flytekit.types.iterator.iterator import IteratorTransformer
410

511

612
@task
@@ -21,3 +27,12 @@ def wf(a: int) -> typing.List[int]:
2127

2228
def test_iterator():
2329
assert wf(a=4) == [0, 1, 2, 3]
30+
31+
32+
def test_to_python_value_non_collection_raises_with_message():
33+
ctx = FlyteContextManager.current_context()
34+
lit = Literal(scalar=Scalar(primitive=Primitive(integer=42)))
35+
trans = IteratorTransformer()
36+
with pytest.raises(TypeTransformerFailedError) as exc_info:
37+
trans.to_python_value(ctx, lit, typing.Iterator[int])
38+
assert exc_info.value.args[0] # args[0] must be set — not empty tuple

0 commit comments

Comments
 (0)