diff --git a/src/_bentoml_sdk/io_models.py b/src/_bentoml_sdk/io_models.py index 50f898fa9a9..3bacfc86b7d 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] + iterator_args = get_args(return_annotation) + return_annotation = iterator_args[0] if iterator_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..75ff3405ef2 --- /dev/null +++ b/tests/unit/_bentoml_sdk/test_io_models.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import collections.abc as cabc +import typing as t + +import pytest + +from _bentoml_sdk.io_models import IODescriptor + + +@pytest.mark.parametrize( + "return_annotation", + [ + t.Iterator, + t.Generator, + t.AsyncIterator, + t.AsyncGenerator, + cabc.Iterator, + cabc.Generator, + cabc.AsyncIterator, + cabc.AsyncGenerator, + ], +) +def test_from_output_accepts_bare_iterators(return_annotation: t.Any) -> None: + def stream() -> t.Iterator[t.Any]: + yield None + + stream.__annotations__["return"] = return_annotation + + descriptor = IODescriptor.from_output(stream) + + assert descriptor.model_fields["root"].annotation is t.Any