How can I set up a wrapping decorator for

@op

so that return types are correctly detected? If I declare an op like this:

@op
def my_op(param: str) -> str:
   return f"hello {param}"

Then in the DAG view

my_op

shows up with

(param: str) => (result: str)

, which is great / what I'd expect. I've defined a wrapper for

@op

which does some extra work (this is

@otel_op

in formenergy-observability) using

functools.wraps

, roughly like this:

def otel_op(
    *decorator_args,
    **decorator_kwargs,
):
    """A wrapper for @op which does extra work"""
    def wrapper(func):
        @op(*decorator_args, **decorator_kwargs)
        @wraps(func)
        def raw_op(context, *func_args, **func_kwargs):
            with my_custom_context_manager():
                ret = func(context, *func_args, **func_kwargs)
                if inspect.isgenerator(ret):
                    # Yield from the generator here rather than just returning it so
                    # that the generator executes within the context manager.
                    yield from ret
                else:
                    yield Output(ret, "result")
        return raw_op

    return wrapper

@otel_op()
def my_op(param: str) -> str:
   return f"hello {param}"

With this, in the DAG view

my_op

shows up with the right parameter type, but no return type detected:

(param: str) => (result: Any)

. Any pointers for why it wouldn't be detecting my

-> str

any more? If I explicitly declare

@otel_op(out=Out(str))
def my_op(param: str):
   return f"hello {param}"

then it does show

(param: str) => (result: str)

, but I'd rather stick to plain Python return type annotations where I can.

You might need to use the

out

parameter in the

op

decorator to specify specific output types. So like

@op(
   out=Out(str)
   ...
)