#265 Fix mypy tasks Pull request opened by cosmicBboy This PR adds the

F

function type to the

TaskTemplate

generic type def, and `@overload`s to the

@env.task

method, task template

__call__

and

aio

methods to make the mypy linter happy. Without it, it tries to type check operations between

Coroutine[Any, Any, R]

with types in user code, resulting in mypy errors. Example: """Basic task type checking test.""" from flyte import TaskEnvironment env = TaskEnvironment(name="test_env") @env.task async def simple_task(x: int) -> int: """Simple async task with proper types.""" return x * 2 @env.task def sync_task(x: str) -> str: """Simple sync task with proper types.""" return x.upper() @env.task async def multi_input_task(x: int, y: str, z: float = 1.0) -> tuple[int, str, float]: """Task with multiple inputs and tuple return.""" a = await simple_task(x) b = await sync_task.aio(y) o1 = a + 1 o2 = b + "foo" o3 = z + 1.0 return o1, o2, o3 mypy error:

tests/mypy/modules/basic_task.py:24: error: Item "function" of "AsyncFunctionTaskTemplate[[str], str] | Callable[[str], str]" has no attribute "aio"  [union-attr]
tests/mypy/modules/basic_task.py:25: error: Unsupported operand types for + ("Coroutine[Any, Any, int]" and "int")  [operator]
tests/mypy/modules/basic_task.py:25: note: Maybe you forgot to use "await"?
tests/mypy/modules/basic_task.py:25: note: Left operand is of type "Coroutine[Any, Any, int] | int"
tests/mypy/modules/basic_task.py:26: error: Unsupported left operand type for + ("Coroutine[Any, Any, str]")  [operator]
tests/mypy/modules/basic_task.py:26: note: Maybe you forgot to use "await"?
tests/mypy/modules/basic_task.py:26: note: Left operand is of type "Coroutine[Any, Any, str] | str | Any"
Found 3 errors in 1 file (checked 1 source file)

flyteorg/flyte-sdk