diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index ebf1921871..e671347cee 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -3,6 +3,7 @@ import collections import copy import inspect +import sys import typing from collections import OrderedDict from typing import Any, Dict, Generator, List, Optional, Tuple, Type, TypeVar, Union, cast @@ -381,10 +382,12 @@ def transform_function_to_interface(fn: typing.Callable, docstring: Optional[Doc return_annotation = type_hints.get("return", None) ctx = FlyteContextManager.current_context() - # Only check if the task/workflow has a return statement at compile time locally. if ( ctx.execution_state + # Only check if the task/workflow has a return statement at compile time locally. and ctx.execution_state.mode is None + # inspect module does not work correctly with Python <3.10.10. https://github.com/flyteorg/flyte/issues/5608 + and sys.version_info >= (3, 10, 10) and return_annotation and type(None) not in get_args(return_annotation) and return_annotation is not type(None) diff --git a/tests/flytekit/unit/core/test_workflows.py b/tests/flytekit/unit/core/test_workflows.py index 43635bcbbb..8cb72aadec 100644 --- a/tests/flytekit/unit/core/test_workflows.py +++ b/tests/flytekit/unit/core/test_workflows.py @@ -246,6 +246,29 @@ def one_output_wf() -> int: # type: ignore one_output_wf() +def test_custom_wrapper(): + def our_task( + _task_function: typing.Optional[typing.Callable] = None, + **kwargs, + ): + def wrapped(_func: typing.Callable): + return task(_task_function=_func) + + if _task_function: + return wrapped(_task_function) + else: + return wrapped + + @our_task( + foo={ + "bar1": lambda x: print(x), + "bar2": lambda x: print(x), + }, + ) + def missing_func_body() -> str: + return "foo" + + def test_wf_no_output(): @task def t1(a: int) -> int: