Improve task decorator type hints with overload

[ringohoffman]

TL;DR

Functions decorated by flytekit.task are not hinted as PythonFunctionTask as a result of its return typing, so there are spurious type errors when attempting to register such tasks on FlyteRemote objects.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

Here you can see an explanation of the problem and the fix:

import flytekit
import flytekit.configuration as flyte_config
import flytekit.remote


@flytekit.task
def my_task1() -> int:
    return 0


@flytekit.task()
def my_task2() -> int:
    return 0


@flytekit.task(task_config=5)
def my_task3() -> int:
    return 0


my_task4 = flytekit.task(lambda: 0)

remote = flytekit.remote.FlyteRemote(...)  # type: ignore
serialization_settings = flyte_config.SerializationSettings(...)  # type: ignore

# before
reveal_type(my_task1)  # Type of "my_task1" is "((...) -> Any) | PythonFunctionTask"
remote.register_task(my_task1, serialization_settings)  # error: Argument of type "((...) -> Any) | PythonFunctionTask" cannot be assigned to parameter "entity" of type "PythonTask" in function "register_task"

reveal_type(my_task2)  # Type of "my_task2" is "() -> int"
remote.register_task(my_task2, serialization_settings)  # error: Argument of type "() -> int" cannot be assigned to parameter "entity" of type "PythonTask" in function "register_task"

reveal_type(my_task3)  # Type of "my_task3" is "() -> int"
remote.register_task(my_task3, serialization_settings)  # error: Argument of type "() -> int" cannot be assigned to parameter "entity" of type "PythonTask" in function "register_task"

reveal_type(my_task4)  # Type of "my_task4" is "((...) -> Any) | PythonFunctionTask"
remote.register_task(my_task4, serialization_settings)  # Argument of type "((...) -> Any) | PythonFunctionTask" cannot be assigned to parameter "entity" of type "PythonTask" in function "register_task"


# after
reveal_type(my_task1)  # Type of "my_task1" is "PythonFunctionTask"
remote.register_task(my_task1, serialization_settings)  # OK

reveal_type(my_task2)  # Type of "my_task2" is "PythonFunctionTask[T@task]"
remote.register_task(my_task2, serialization_settings)  # OK

reveal_type(my_task3)  # Type of "my_task3" is "PythonFunctionTask[int]"
remote.register_task(my_task3, serialization_settings)  # OK

reveal_type(my_task4)  # Type of "my_task4" is "PythonFunctionTask"
remote.register_task(my_task4, serialization_settings)  # OK

I accomplished this fix by using typing.overload to hint a different return type based on whether _task_function is passed to flytekit.task or not.

If nothing is passed as in the case of my_task2 or my_task3, we hint the return type as Callable[[Callable[..., Any]], PythonFunctionTask[T]]... that is a callable that accepts a single argument that is any function and returns a PythonFunctionTask.

If something is passed for _task_function as in my_task1 or my_task4, we hint the return type as PythonFunctionTask[T].

Tracking Issue

N/A

Follow-up issue

N/A

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).

[codecov]

Codecov Report

Merging #1631 (95975e2) into master (e44b802) will decrease coverage by 0.02%.
The diff coverage is 11.11%.

@@            Coverage Diff             @@
##           master    #1631      +/-   ##
==========================================
- Coverage   71.22%   71.20%   -0.02%     
==========================================
  Files         334      334              
  Lines       30391    30398       +7     
  Branches     5490     5492       +2     
==========================================
  Hits        21645    21645              
- Misses       8206     8213       +7     
  Partials      540      540              

Impacted Files	Coverage Δ
flytekit/core/task.py	28.84% <11.11%> (-4.49%)	⬇️

[welcome]

Congrats on merging your first pull request! 🎉

[wild-endeavor]

@ringohoffman thanks very much for the contribution!

[wild-endeavor]

Ooh @ringohoffman do we need to do something similar to this for workflow?

[ringohoffman]

Ooh @ringohoffman do we need to do something similar to this for workflow?

It looks like it would be more correct to do it similarly to the way this one was done. It's not true that flytekit.core.workflow.workflow always returns a WorkflowBase. I can open a PR for it.

[pingsutw]

@ringohoffman please do it, thanks!