Flyte enables production-grade orchestration for machine learning workflows and data processing created to accelerate local workflows to production.

Flyte

image.png

:wave:  I have a couple tasks and a workflow that calls those. The issue I’m facing is that since `flytekit.task` returns a union of `PythonFunctionTask`, `FuncOut`, `Promise` the typechecker’s utility is greatly reduced.
Is there a paved path to preserving the original function’s return type? Or some other way that is peaceful with the typechecker (using pyright)?

Code:
```@dataclass_json
@dataclass
class InnerDC:
    subfield: str
    subpath: str


@dataclass_json
@dataclass
class OuterDC:
    name: str
    path: str
    inner_dc: InnerDC
    uid: Optional[int] = None


@task
def make_test_dc(name: str, path: str, subfield: str, subpath: str) -&gt; OuterDC:
    return OuterDC(name=name, path=path, inner_dc=InnerDC(subfield=subfield, subpath=subpath))


@task
def extract_subpath(dc: OuterDC) -&gt; str:
    return dc.inner_dc.subpath


@workflow
def wf(name: str, path: str, subfield: str, subpath: str) -&gt; str:
    dc = make_test_dc(name=name, path=path, subfield=subfield, subpath=subpath)
    return extract_subpath(dc=dc)```

this is a known issue though, and something we were at least planning on working on… soon

I see, thanks for the info. Just out of curiosity, what would the solution look like?
Currently I am using a `TypedTask` protocol to overload the return type to be the same as the decorated function which is a decent compromise for now, but with some obvious downsides (no kwargs info, no docstrings etc)

the solution has yet to be discovered tbh.  we were gonna start thinking about it next week tbh.  at least planning to.

but it’d likely involve exporting a series of pyi files along with flytekit that can hopefully teach lsps how to interpret flytekit code correctly.

That sounds amazing! Much appreciated :pray:

<https://github.com/flyteorg/flyte/issues/3681> <https://github.com/flyteorg/flyte/issues/2147> <https://github.com/flyteorg/flyte/issues/3682> <https://github.com/flyteorg/flyte/issues/2525>

all the issues we’ve collected.  feel free to upvote one or all of them.  this is something we definitely need to work on

Definitely more of a nice-to-have, but it is a bit ironic how typing is at the core of flyte, but we don’t get good typing support with mypy, etc at the workflow and task definitions.

rest assured you're not the only one, <@U05AJR5NDB4> . We're actively looking for solutions. :slightly_smiling_face:

I opened <https://github.com/flyteorg/flytekit/pull/2411> to get the ball rolling. I do not think it is the most correct solution, but it does enable type checking.