refactor(core): Enhance return type extraction logic

flytekit/core/interface.py

@@ -507,7 +507,9 @@ def t(a: int, b: str) -> Dict[str, int]: ...
 
     # This statement results in true for typing.Namedtuple, single and void return types, so this
     # handles Options 1, 2. Even though NamedTuple for us is multi-valued, it's a single value for Python
-    if isinstance(return_annotation, Type) or isinstance(return_annotation, TypeVar):  # type: ignore
+    if hasattr(return_annotation, "__bases__") and (
+        isinstance(return_annotation, Type) or isinstance(return_annotation, TypeVar)  # type: ignore
+    ):
         # isinstance / issubclass does not work for Namedtuple.
         # Options 1 and 2
         bases = return_annotation.__bases__  # type: ignore

tests/flytekit/unit/core/test_interface.py

@@ -4,7 +4,7 @@
 from typing import Dict, List
 
 import pytest
-from typing_extensions import Annotated  # type: ignore
+from typing_extensions import Annotated, TypeVar  # type: ignore
 
 from flytekit import map_task, task
 from flytekit.core import context_manager
@@ -96,6 +96,15 @@ def t(a: int, b: str) -> Dict[str, int]:
     assert len(return_type) == 1
     assert return_type["o0"] == Dict[str, int]
 
+    VST = TypeVar("VST")
+
+    def t(a: int, b: str) -> VST:  # type: ignore
+        ...
+
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
+    assert len(return_type) == 1
+    assert return_type["o0"] == VST
+
 
 def test_named_tuples():
     nt1 = typing.NamedTuple("NT1", x_str=str, y_int=int)