ray.serve.batch return type problems with pylance and mypy

[ArthurBook]

What happened + What you expected to happen

Hi! first of all, thank you for the work on this amazing project. I am a huge fan.

There is a type-hint nit that has been bothering me in the ray.serve.batch decorator. It is causing some issues with pylance and mypy.

The issue stems from that the bounds on the typevars F and G have TypeVar bound types, which aren't supported. Another related issue is the behavior of decorating a method on a class with @batch. Since the method has an implicit self argument, the Callable typehint is unable to match the signature.

In other words, the return-type for a @batch decorated function is not inferred correctly, making mypy / pylance (probably other checkers too) lost the function signature for the decorated function/method.

Versions / Dependencies

poetry show ray name : ray
version : 2.11.0
description : Ray provides a simple, universal API for building distributed applications.

dependencies

• aiohttp >=3.7
• aiohttp-cors *
• aiosignal *
• click >=7.0
• colorful *
• fastapi *
• filelock *
• frozenlist *
• fsspec *
• grpcio >=1.32.0
• grpcio >=1.42.0
• jsonschema *
• memray *
• msgpack >=1.0.0,<2.0.0
• numpy >=1.20
• opencensus *
• packaging *
• pandas >=1.3
• pandas *
• prometheus-client >=0.7.1
• protobuf >=3.15.3,<3.19.5 || >3.19.5
• py-spy >=0.2.0
• pyarrow >=6.0.1
• pydantic <2.0.dev0 || >=2.5.dev0,<3
• pyyaml *
• requests *
• smart-open *
• starlette *
• tensorboardX >=1.9
• uvicorn *
• virtualenv >=20.0.24,<20.21.1 || >20.21.1
• watchfiles *

Reproduction script

import ray.serve

@ray.serve.batch(max_batch_size=10)
def batch_fn(integers: list[int]) -> list[int]:
    return [i * 2 for i in integers]

pylance (vscode):

Argument of type "(integers: list[int]) -> list[int]" cannot be assigned to parameter of type "F@batch"
  Type "(integers: list[int]) -> list[int]" cannot be assigned to type "(List[T]) -> List[R]"
    Type "(integers: list[int]) -> list[int]" cannot be assigned to type "(List[T]) -> List[R]"
      Parameter 1: type "List[T]" cannot be assigned to type "list[int]"
        "List[T]" is incompatible with "list[int]"
          Type parameter "_T@list" is invariant, but "T" is not the same as "int"
          Consider switching from "list" to "Sequence" which is covariantPylancereportArgumentType

mypy:

poetry run mypy ray_batch_type.py 
labeling/ray_batch_type.py:4: error: Value of type variable "F" of function cannot be "Callable[[list[int]], list[int]]"  [type-var]
labeling/ray_batch_type.py:9: error: Never not callable  [misc]

Issue Severity

Low: It annoys or frustrates me.