@weave.op decorator destroys type information

[pokey]

If I have the following code:

@weave.op()
def get_first_draft(
    previous_episodes: list[EpisodeCore],
    repository_data: dict,
) -> str:

VSCode type checking stops checking for valid args when I call the function, and autocomplete breaks for the args. The function is just interpreted as an Op object. See eg the hover

This is a showstopper, as I rely heavily on autocomplete and type checking. Notice how I forgot to add an arg

[pokey]

For now, this is my workaround:

from collections.abc import Callable
from typing import TypeVar, cast

import weave

C = TypeVar("C", bound=Callable)


def weave_op():
    def ret(fn: C) -> C:
        return cast(C, weave.op()(fn))

    return ret

Which I then use like:

@weave_op()
def my_func(foo: str) -> str:
    ...

I don't support any of the @weave.op() arguments, and I don't expose any typing information about what I can use on an Op, but it's good enough for me for now. I think the correct thing would be to make the Op class have a generic type argument

[jamie-rasmussen]

Nice workaround! Glad you are unblocked. Will open an internal ticket to track this.