#1635 Improve workflow decorator type hints with overload Pull request opened by ringohoffman TL;DR Currently, the

workflow

decorator is hinted as always returning a

WorkflowBase

, which is not true when

_workflow_function

is `None`; this leads to a spurious type error depending on how

workflow

is called; similar to #1631, I propose using

typing.overload

to differentiate the return type of

workflow

based on the value of

_workflow_function

. 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 is an example of the typing bug this fix addresses:

import flytekit
import flytekit.remote


@flytekit.workflow
def my_workflow1() -> int:
    return 0


@flytekit.workflow()
def my_workflow2() -> int:
    return 0


my_workflow3 = flytekit.workflow(lambda: 0)

remote = flytekit.remote.FlyteRemote(...)  # type: ignore

# before
reveal_type(my_workflow1)  # Type of "my_workflow1" is "WorkflowBase"
remote.register_workflow(my_workflow1)  # OK

reveal_type(my_workflow2)  # Type of "my_workflow2" is "Tuple[Promise] | Promise | VoidPromise | Tuple | None"
remote.register_workflow(my_workflow2)  # error: Argument of type "Tuple[Promise] | Promise | VoidPromise | Tuple | None" cannot be assigned to parameter "entity" of type "WorkflowBase" in function "register_workflow"

reveal_type(my_workflow3)  # Type of "my_workflow3" is "WorkflowBase"
remote.register_workflow(my_workflow3)  # OK


# after
reveal_type(my_workflow1)  # Type of "my_workflow1" is "PythonFunctionWorkflow"
remote.register_workflow(my_workflow1)  # OK

reveal_type(my_workflow2)  # Type of "my_workflow2" is "PythonFunctionWorkflow"
remote.register_workflow(my_workflow2)  # OK

reveal_type(my_workflow3)  # Type of "my_workflow3" is "PythonFunctionWorkflow"
remote.register_workflow(my_workflow3)  # OK

I am proposing to change the return type from

WorkflowBase

to

PythonFunctionWorkflow

. This is the class being instantiated and seems to match what is done by

task

, which returns

PythonFunctionTask

instead of the base class

PythonTask

which is expected by

FlyteRemote.register_task

. I am also proposing we use the condition

if _workflow_function is not None:

instead of

if _workflow_function:

. This is a recommendation by the Google Python Style Guide: https://google.github.io/styleguide/pyguide.html#2144-decision

Always use if foo is None: (or is not None) to check for a None value. E.g., when testing whether a variable or argument that defaults to None was set to some other value. The other value might be a value that’s false in a boolean context!

Tracking Issue N/A Follow-up issue N/A flyteorg/flytekit Codecov: 30.00% of diff hit (target 71.24%) ✅ 29 other checks have passed 29/30 successful checks