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

Flyte

Is there support for tasks with inputs that use Unions? `Union[str, flytekit.types.directory.FlyteDirectory]`. Trying to create a task that could take in a str or a flytedirectory but I get

```    raise ValueError(f"Expected a directory. {source_path} is not a directory")
ValueError: Expected a directory. X is not a directory

AssertionError: Failed to Bind variable input for function ```

Yea it does to seem work; I tested it with a small task, I think for this task its not working bc I have a decorator around it to do some custom io and calling `func(*args, **kwargs)` inside the decorator might be breaking something?

You have to ensure you pass the type info
Correctly 

```def mount(job_type):
    def inner(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):```

I am using the the `@wraps` , seems like I need to use partial. I'll play around with it

I think there might be something special with flytedirectory and decorators? I'm not really sure but here's a minimal example that is failing

```from functools import partial, wraps
def validate_output(fn=None, *, floor=0):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        out = fn(*args, **kwargs)
        if out &lt;= floor:
            raise ValueError(
                f"output of task {fn.__name__} must be a positive number, found {out}"
            )
        return out

    if fn is None:
        return partial(validate_output, floor=floor)

    return wrapper

@task
@validate_output(floor=0)
def test_task(input: Union[str, flytekit.types.directory.FlyteDirectory]) -&gt;int :
    print(input)
    return 5

if __name__ == "__main__":
    result = test_task(
        input="some str",
    )```

I just copied the decorator from the docs and it's still trying to resolve the input as a directory instead of as a string

<@U047KD38XEU> Sorry, there is a small bug in type engine. I just fixed in this <https://github.com/flyteorg/flytekit/pull/1270|PR>. you can try this by running `pip install git+<https://github.com/flyteorg/flytekit@flyte_dir_union>`

yes, it seems to work with the code snippet I shared above! thanks so much :heart: