From 090905235bbcadb48ca23a1fda59decb3dc594a2 Mon Sep 17 00:00:00 2001 From: Bernhard Stadlbauer Date: Wed, 15 Dec 2021 11:37:44 +0100 Subject: [PATCH 1/4] When using the `task` and `workflow` decorator, correctly wrap the function This enables tooling such as docstring search tools to unwrap the object and show the correct docstring. Signed-off-by: Bernhard Stadlbauer --- flytekit/core/task.py | 3 ++- flytekit/core/workflow.py | 2 ++ tests/flytekit/unit/core/test_wrapping.py | 18 ++++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) create mode 100644 tests/flytekit/unit/core/test_wrapping.py 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/tests/flytekit/unit/core/test_wrapping.py b/tests/flytekit/unit/core/test_wrapping.py new file mode 100644 index 0000000000..1102c682ed --- /dev/null +++ b/tests/flytekit/unit/core/test_wrapping.py @@ -0,0 +1,18 @@ +# Copyright (C) 2015-2021 Blackshark.ai GmbH. All Rights reserved. www.blackshark.ai +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_wf_correctly_wrapped(): + @workflow + def my_workflow(a: int) -> int: + return a + + assert my_workflow.__wrapped__ == my_workflow._workflow_function From 5f2b87f0f3279f849c9ebfe3a0367d01089c2065 Mon Sep 17 00:00:00 2001 From: Bernhard Stadlbauer Date: Thu, 16 Dec 2021 16:16:17 +0100 Subject: [PATCH 2/4] Remove blackshark copyright header Signed-off-by: Bernhard Stadlbauer --- tests/flytekit/unit/core/test_wrapping.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_wrapping.py b/tests/flytekit/unit/core/test_wrapping.py index 1102c682ed..12cd581e66 100644 --- a/tests/flytekit/unit/core/test_wrapping.py +++ b/tests/flytekit/unit/core/test_wrapping.py @@ -1,4 +1,3 @@ -# Copyright (C) 2015-2021 Blackshark.ai GmbH. All Rights reserved. www.blackshark.ai from flytekit import task, workflow From 0350598fcf058ac1e29dbe597def68d2f222e512 Mon Sep 17 00:00:00 2001 From: Bernhard Stadlbauer Date: Thu, 16 Dec 2021 16:20:51 +0100 Subject: [PATCH 3/4] Fix broken great expectations test Signed-off-by: Bernhard Stadlbauer --- plugins/flytekit-greatexpectations/tests/test_task.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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(): From 43bf75b8ffdf216cd873183c614d81808121067b Mon Sep 17 00:00:00 2001 From: Bernhard Stadlbauer Date: Fri, 17 Dec 2021 07:37:05 +0100 Subject: [PATCH 4/4] Add test for stacked decorators Signed-off-by: Bernhard Stadlbauer --- tests/flytekit/unit/core/test_wrapping.py | 41 +++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/flytekit/unit/core/test_wrapping.py b/tests/flytekit/unit/core/test_wrapping.py index 12cd581e66..97ff2bafec 100644 --- a/tests/flytekit/unit/core/test_wrapping.py +++ b/tests/flytekit/unit/core/test_wrapping.py @@ -1,3 +1,5 @@ +from functools import wraps + from flytekit import task, workflow @@ -9,6 +11,45 @@ def my_task(a: int) -> int: 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: