import ray
ray.init()
@ray.remote
def do_things(x: int, y: float):
 return x * y
@ray.remote
def do_more_things(name: str, value: float):
 return f"{name}: {value}"
pure_value = do_things(x=3, y=2.2) # pure_value is type: float
ref_value = do_things.remote(x=3, y=2.2) # ref_value is type: ObjectRef[float]
pure_second_value = do_more_things(name="Foo", value=pure_value) # this is fine
ref_second_value = do_more_things.remote(name="Foo", value=ref_value) # this should be fine

do_things.remote(x: int | ObjectRef[int], y: float | ObjectRef[float]) -> ObjectRef[float]: ...
do_more_things.remote(name: str | ObjectRef[str], value: float | ObjectRef[float]) -> ObjectRef[str]: ...

# do_things doesn't/shouldn't support y: float | ObjectRef[float], only y: float
do_things(x=3, y=ref_value)
# do_more_things.remote doesn't/shouldn't support name: str| ObjectRef[float], only name: str | ObjectRef[str]
do_more_things.remote(name=ref_value, value=2.2)

from typing import Callable, TypeVar
from typing_extensions import ParamSpec
from ray import ObjectRef
P = ParamSpec("P")
R = TypeVar("R")
def remotify(func: Callable[P, R]) -> Callable[P, ObjectRef[R]]:
 # Do the equivalent of @ray.remote here to wrap/modify the function
 return func
def do_things(x: int, y: float):
 return x * y
remote_do_things = remotify(do_things)
def do_more_things(name: str, value: float):
 return f"{name}: {value}"
remote_do_more_things = remotify(do_more_things)
# Now this has valid types, inline error checks, autocompletion
pure_value = do_things(x=3, y=2.2)
# Now this has *almost* valid types, inline error checks, autocompletion
ref_value = remote_do_things(x=3, y=2.2) # ref_value is type: ObjectRef[float]
# But the type annotations for the arguments don't accept the altrenative OjectRef[X]
# This is fine, and typed correctly
pure_second_value = do_more_things(name="Foo", value=pure_value)
# This is fine in code, and the type information is *almost* fine, it should accept
# value: float | ObjectRef[float]
# But it can only take the same signature as the original, with value: float
ref_second_value = remote_do_more_things(name="Foo", value=ref_value)

remotify(do_things)(x=3, y=2.2)

from typing import Callable, TypeVar, Union, overload
from ray import ObjectRef
T0 = TypeVar("T0")
T1 = TypeVar("T1")
T2 = TypeVar("T2")
R = TypeVar("R")
@overload
def remotify(func: Callable[[T0], R]) -> Callable[[Union[T0, ObjectRef[T0]]], ObjectRef[R]]: ...
@overload
def remotify(func: Callable[[T0, T1], R]) -> Callable[[Union[T0, ObjectRef[T0]], Union[T1, ObjectRef[T1]]], ObjectRef[R]]: ...
@overload
def remotify(func: Callable[[T0, T1, T2], R]) -> Callable[[Union[T0, ObjectRef[T0]], Union[T1, ObjectRef[T1]], Union[T2, ObjectRef[T2]]], ObjectRef[R]]: ...
def remotify(func: Callable[..., R]) -> Callable[..., ObjectRef[R]]:
 # Do the equivalent of @ray.remote here to wrap/modify the function
 return func
def do_things(x: int, y: float):
 return x * y
remote_do_things = remotify(do_things)
def do_more_things(name: str, value: float):
 return f"{name}: {value}"
remote_do_more_things = remotify(do_more_things)
# This works
pure_value = do_things(x=3, y=2.2)
# Now this has valid types, but doesn't accept keyword arguments
ref_value = remote_do_things(x=3, y=2.2) # ref_value is type: ObjectRef[float]
# This is fine, and typed correctly
pure_second_value = do_more_things(name="Foo", value=pure_value)
# The type information is fine, but it doesn't accept keyword arguments
ref_second_value = remote_do_more_things(name="Foo", value=ref_value)

from typing import Callable, TypeVar
from typing_extensions import TypeVarTuple, Unpack
from ray import ObjectRef
T = TypeVarTuple("T")
R = TypeVar("R")
def remotify(func: Callable[[Unpack[T]], R]) -> Callable[[Unpack[T]], ObjectRef[R]]:
 # Do the equivalent of @ray.remote here to wrap/modify the function
 return func