Adds typing to @task and @workflow

[thomasjpfan]

Tracking issue

Closes flyteorg/flyte#3681

Why are the changes needed?

Type checkers do not work with flytekit right now.

What changes were proposed in this pull request?

With this PR, the @task and @workflow pretends to return their original types for type checking purposes. The typing in this PR is not strictly correct, because the workflow transforms everything into Promises. I think this PR is an okay solution to get typing to work for user's workflows.

As for actually typing the Promises, I do not see a way to actually type them. Python typing allows us to change the output of the decorator, but we can not change the inputs:

def task_example(f: Callable[P, R]) -> Callable[P, Promise[R]]:
    # We need to somehow express that the inputs are also Promises, and not the original types.
    def inner(*args: P.args, **kwargs: P.kwargs) -> Promise[R]:
        return f(*args, **kwargs)

    return inner

How was this patch tested?

Run mypy on workflow code and see failures when types do not match:

from flytekit import task, workflow


@task(cache=False, cache_version="v1")
def add(a: int, b: int) -> int:
    return a + b


# note how `b` is the incorrect type
@workflow
def wf(a: int, b: str) -> int:
    return add(a=a, b=b)

Running mypy outputs:

error: Argument "b" to "add" has incompatible type "str"; expected "int"  [arg-type]
Found 1 error in 1 file (checked 1 source file)

[eapolinario]

@thomasjpfan , I like this as this fixes the problem for @task and @workflow decorators, but we can't simply enable py.typed, right? Otherwise all entities exported by flytekit will be type-checkable.

[thomasjpfan]

I like this as this fixes the problem for @task and @workflow decorators, but we can't simply enable py.typed, right?

If we remove the py.typed, then mypy will continue to skip @task and @workflow`:

file.py:1: error: Cannot find implementation or library stub for module named "flytekit"  [import-not-found]
file.py:1: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
Found 1 error in 1 file (checked 1 source file)

Users will have to wait until we add py.typed before they can type check their @task and @workflow.

In any case, I am okay with removing the flag, if we are okay with users waiting.

[wild-endeavor]

I think users are okay waiting. there's just too much in flytekit itself to clean up. This will be an ongoing effort. In any case this will make pyright less red at least. Let's remove and merge? Is the lint error caused by the py.typed file itself?

[eapolinario]

The rules for distributing type information are described here. We might need to play with adding *.pyi files alongside the py.typed marker to find the right combination.