Wrong type hint for return types from a task

[gorootde]

Bug summary

Any task's return type hint is lost when annotating it with @task.

Example 1:

class Foo:
    def foo(self):
        pass   

@task
def get_foo() -> Foo:
  return Foo()

@task
def get_list() -> List[Any]:
  ...

@flow
def run():
   result = get_foo() #Type of result is None and not Foo
   result2 = get_list() # Type of result2 is None and not List[Any]

Version info

Version:             3.1.10
API version:         0.8.4
Python version:      3.11.10
Git commit:          b11b56fc
Built:               Tue, Dec 24, 2024 10:58 PM
OS/Arch:             darwin/arm64
Profile:             local
Server type:         server
Pydantic version:    2.10.4

Additional context

No response

[zzstoatzz]

hi @gorootde - can you say more about how you're seeing a loss of type information? ie what is reporting this?

• I am not able to reproduce the problem with get_foo
• get_list is returning None because you have ... in place of the implementation

from typing import Any, reveal_type

from prefect import flow, task


class Foo:
    def foo(self):
        pass


@task
def get_foo() -> Foo:
    return Foo()


@task
def get_list() -> list[Any]: ...


@task
def actually_get_list() -> list[Any]:
    return [1, 2, 3]


@flow
def run():
    result = get_foo()
    assert isinstance(result, Foo)
    result2 = get_list()
    assert result2 is None
    result3 = actually_get_list()
    assert isinstance(result3, list)
    reveal_type(result)  # Runtime type is 'Foo'
    reveal_type(result2)  # Runtime type is 'NoneType'
    reveal_type(result3)  # Runtime type is 'list'


if __name__ == "__main__":
    run()