How to specify ParamSpec signature in assert_type?

[randolf-scholz]

I am writing some type-hints for pytorch. The jit.script functions can be used to create compiled versions of functions.

P = ParamSpec("P")

@overload
def script(  # type: ignore[misc]
    obj: Callable[P, R],
    optimize: Optional[bool] = None,
    _frames_up: int = 0,
    _rcb: Optional[ResolutionCallback] = None,
    example_inputs: Union[List[Tuple], Dict[Callable, List[Tuple]], None] = None,
) -> ScriptFunction[P, R]:

and

class ScriptFunction(Callable[P, R]):
    def __call__(self, *args: P.Args, **kwargs: P.kwargs) -> R: ...

Now, I want to add a test:

assert_type(jit.script(nn.functional.relu), Callable[[Tensor, bool], Tensor])
assert_type(jit.script(nn.functional.relu), ScriptFunction[[Tensor, bool], Tensor])

which results in

test/typing/pass/jit.py:28:1: error: Expression is of type "ScriptFunction[[Tensor, bool], Tensor]", not "Callable[[Tensor, bool], Tensor]"  [assert-type]
test/typing/pass/jit.py:29:1: error: Expression is of type "ScriptFunction[[Arg(Tensor, 'input'), DefaultArg(bool, 'inplace')], Tensor]", not "ScriptFunction[[Tensor, bool], Tensor]"  [assert-type]
Found 2 errors in 1 file (checked 3 source files)

How can I specify the correct signature in assert_type?

[randolf-scholz]

The best I could come up with is

relu: Callable[[Tensor, bool], Tensor] = nn.functional.relu  # forget argument names
assert_type(jit.script(relu), ScriptFunction[[Tensor, bool], Tensor])

There really ought to be a way to reference the signature of another callable, like:

relu_siganture = signature[relu.__call__]
assert_type(jit.script(relu), ScriptFunction[relu_signature.params, relu_signature.return])

[Daverball]

The point of assert_type is to be as strict as possible and only allow exact matches, your first statement already sort of does the check you want to make, since a static type checker would complain if the signature wasn't compatible with the positional only version.

I agree that it would be nice to have a way to specify all the different kinds of arguments in a ParamSpec and mypy has an old unsupported representation of those arguments, however it will not allow them to be used in ParamSpec, even if it still sort of generates them and uses them internally.

I think the following is perfectly fine for testing what you wanted to test. It won't catch Any, so it's not a perfect replacement for assert_type on the exact ParamSpec, but it beats duplicating the signature of relu just so you can verify the signature of relu, which is kind of a nonsensical test. You're kind of just verifying whether a type checker is implementing ParamSpec correctly at that point, which should be part of the type checkers internal test suite.

relu: ScriptFunction[[Tensor, bool], Tensor] = jit.script(nn.functional.relu)