How to annotate modifications for function arguments with ParamSpec and TypeVar

[tiangolo]

I want to be able to annotate something that takes a function and creates a new one that alters the argument types it accepts.

A quick example, using Ray:

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

So, do_things() and do_more_things() should keep their signature. But their remote counterparts should have as signature that allows the original types or a wrapper ObjectRef[OriginalType] for each argument:

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]: ...

Now, I know type annotating this in this particular API design is probably difficult, and possibly not easily supported. I was thinking of an alternative API that could take advantage of ParamSpec, TypeVars with overloads and/or TypeVarTuple. But I still can't find a way to get all the desired features together.

Wanted features

• I would want to have editor inline errors for invalid types (mypy, Pyright), for example, these would be errors:

# 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)

• I would want the correct usage (as the first example) to be considered correct, with the correct/valid type information.
• I would want to have autocompletion in editors for the modified (remote) counterpart, for keyword arguments.
• I would want the modified (remote) version to accept keyword arguments and validate its types, with the extended/modified type accepting ObjectRef[X].

Alternatives

I have been playing around with alternative code interfaces/APIs to achieve this.

ParamSpec

For example with ParamSpec and a wrapper function that doesn't create a new attribute:

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)

And with this approach it could be inlined, which is what I would expect would be the common usage pattern, this allows type checks (inline errors) and autocompletion in editors:

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

But this approach doesn't support extending the types of the arguments with ObjectRef[X].

TypeVars and overloads

Another alternative is with TypeVars and overload, this enables type checks, extending the types, but it no longer supports keyword arguments (nor autocompletion for them in editors):

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)

TypeVarTuple

I also tried experimenting with TypeVarTuple:

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

But the end result is both of the two problems above combined 😔 , I can't update the new type arguments in the resulting function and it doesn't support keyword arguments.

Questions

I would like to have both of these main features combined:

• Having autocompletion for keyword arguments AND
• Being able to extend the received types with their alternative ObjectRef version

Both of the two first alternatives achieve almost all of the things I would want. I also tried using overloads combining both ideas, but only one of the signatures would be taken into account. So there's always one of these two main features missing.

---

I'm not even sure the title of the discussion is right or what is the right term for this. I've been trying to find the solution for a couple of days in the discussions and issues here, the typing mailing list, the Pyright discussions, etc. But I didn't find any previous discussions around it.

Is there any way to achieve annotating types in some way that solve this problem? Am I missing something else?

Edit 2022-04-29

Assuming this is currently not possible, what would be needed to make it possible? Would some way of achieving this be acceptable?

And if so, what would be the best approach to make it possible?

I'm not sure what's the process, but maybe there could be a way to sponsor someone with the right expertise here to tackle it, I imagine it would require a PEP, work on mypy, not sure what else, but maybe I'm being naive in some way and it would require a different approach.

[hmc-cs-mdrissi]

You are looking for a Map operator on types. An older version of typevartuple pep had map operator and it would have let last approach work ignoring keyword argument support. Something like,

NewT = Union[T, ObjectRef[T]]

Callable[[Unpack[Ts]], R] -> Callable[[Unpack[Map[NewT, Ts]]], R]

Map takes a generic type/alias as first argument and then applies it over each type in typevartuple.

A Map type operator that works on paramspecs in theory could also be proposed.

At moment though the current type system does not support this. My guess is eventually TypeVarTuple Map will be proposed as reason it was dropped was to simplify pep 646 and leave it as a future follow up. I’m unaware of any plan to support Map on paramspec keyword arguments.

[tiangolo]

Ah, that sounds very interesting, at least it seems I'm not omitting something obvious, thank you!

[eapolinario]

@tiangolo , did you end up finding a solution for this problem? I'm facing a similar problem in flytekit (well described in flyteorg/flyte#3682 (comment)).