diff --git a/src/_bentoml_sdk/io_models.py b/src/_bentoml_sdk/io_models.py index 50f898fa9a9..24e89da30c3 100644 --- a/src/_bentoml_sdk/io_models.py +++ b/src/_bentoml_sdk/io_models.py @@ -432,7 +432,8 @@ def from_output(cls, func: t.Callable[..., t.Any]) -> type[IODescriptor]: ) media_type: str | None = None if is_iterator_type(return_annotation): - return_annotation = get_args(return_annotation)[0] + args = get_args(return_annotation) + return_annotation = args[0] if args else t.Any elif is_annotated(return_annotation): content_type = next( (a for a in get_args(return_annotation) if isinstance(a, ContentType)), diff --git a/tests/unit/_bentoml_sdk/test_io_models.py b/tests/unit/_bentoml_sdk/test_io_models.py new file mode 100644 index 00000000000..c14da4dcdad --- /dev/null +++ b/tests/unit/_bentoml_sdk/test_io_models.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import typing as t + +from _bentoml_sdk.io_models import IODescriptor + + +def test_from_output_bare_iterator_does_not_raise(): + """Bare Iterator/Generator with no type args must not raise IndexError.""" + + def fn() -> t.Iterator: + yield 1 + + spec = IODescriptor.from_output(fn) + assert spec is not None