diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 45b7e2dc94..eb393b063d 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -1,4 +1,5 @@ import datetime as _datetime +from functools import update_wrapper from typing import Any, Callable, Dict, List, Optional, Type, Union from flytekit.core.base_task import TaskMetadata, TaskResolverMixin @@ -195,7 +196,7 @@ def wrapper(fn) -> PythonFunctionTask: execution_mode=execution_mode, task_resolver=task_resolver, ) - + update_wrapper(task_instance, fn) return task_instance if _task_function: diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py index 744ecfbb11..ffa6aae934 100644 --- a/flytekit/core/workflow.py +++ b/flytekit/core/workflow.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from enum import Enum +from functools import update_wrapper from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union from flytekit.common import constants as _common_constants @@ -730,6 +731,7 @@ def wrapper(fn): docstring=Docstring(callable_=fn), ) workflow_instance.compile() + update_wrapper(workflow_instance, fn) return workflow_instance if _workflow_function: diff --git a/plugins/flytekit-greatexpectations/tests/test_task.py b/plugins/flytekit-greatexpectations/tests/test_task.py index 9906085db1..578e9ee791 100644 --- a/plugins/flytekit-greatexpectations/tests/test_task.py +++ b/plugins/flytekit-greatexpectations/tests/test_task.py @@ -149,7 +149,6 @@ def valid_wf(dataset: str = "yellow_tripdata_sample_2019-01.csv") -> int: task_object(dataset=dataset) return my_task(csv_file=dataset) - @pytest.mark.xfail(strict=True) @workflow def invalid_wf(dataset: str = "yellow_tripdata_sample_2019-02.csv") -> int: task_object(dataset=dataset) @@ -158,7 +157,8 @@ def invalid_wf(dataset: str = "yellow_tripdata_sample_2019-02.csv") -> int: valid_result = valid_wf() assert valid_result == 10000 - invalid_wf() + with pytest.raises(ValidationError, match=r".*passenger_count -> expect_column_min_to_be_between.*"): + invalid_wf() def test_ge_workflow(): diff --git a/tests/flytekit/unit/core/test_wrapping.py b/tests/flytekit/unit/core/test_wrapping.py new file mode 100644 index 0000000000..97ff2bafec --- /dev/null +++ b/tests/flytekit/unit/core/test_wrapping.py @@ -0,0 +1,58 @@ +from functools import wraps + +from flytekit import task, workflow + + +def test_task_correctly_wrapped(): + @task + def my_task(a: int) -> int: + return a + + assert my_task.__wrapped__ == my_task._task_function + + +def test_stacked_decorators(): + def task_decorator_1(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_1") + return fn(*args, **kwargs) + + return wrapper + + def task_decorator_2(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_2") + return fn(*args, **kwargs) + + return wrapper + + def task_decorator_3(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_3") + return fn(*args, **kwargs) + + return wrapper + + @task + @task_decorator_1 + @task_decorator_2 + @task_decorator_3 + def my_task(x: int) -> int: + """Some function doc""" + print("running my_task") + return x + 1 + + assert my_task.__wrapped__.__doc__ == "Some function doc" + assert my_task.__wrapped__ == my_task._task_function + assert my_task(x=10) == 11 + + +def test_wf_correctly_wrapped(): + @workflow + def my_workflow(a: int) -> int: + return a + + assert my_workflow.__wrapped__ == my_workflow._workflow_function