[Serve][1/n] Improve type annotations for DeploymentHandle, DeploymentResponse, and DeploymentResponseGenerator

[abrarsheikh]

Fixes https://github.com/ray-project/ray/issues/52654

Summary

This PR improves the generic type annotations on Ray Serve's handle classes to enable better IDE support and type inference. It also adds a type checking test file to verify the annotations work correctly.

Changes

Type Annotation Improvements (handle.py)

• Generic return types: Updated methods to return the generic type parameter R instead of Any:

  • DeploymentResponse.result() → R
  • DeploymentResponse.__await__() → Generator[Any, None, R]
  • DeploymentResponseGenerator.__iter__() → Iterator[R]
  • DeploymentResponseGenerator.__next__() → R
  • DeploymentResponseGenerator.__aiter__() → AsyncIterator[R]
  • DeploymentResponseGenerator.__anext__() → R

• mypy compatibility fixes:

  • Added cast() for Future types in sync/async fetch methods to help mypy understand the conditional types
  • Aligned _DeploymentHandleBase.options() signature with DeploymentHandle.options() to fix override error
  • Refactored remote() to return directly from each branch instead of assigning different class types to a variable

Type Checking Tests (check_handle_typing.py)

Added a mypy test file that verifies:

• Generic type R is preserved through result(), __await__, iteration methods
• Generic type T is preserved through options() and method access
• Placeholder tests for future mypy plugin (commented out) that will verify method return type inference

How to Test

    mypy python/ray/serve/tests/typing_files/check_handle_typing.py python/ray/serve/handle.py  \
        --follow-imports=skip \
        --ignore-missing-imports

Future Work

A mypy plugin can be implemented to infer the return type of .remote() based on which deployment method is being called:

handle: DeploymentHandle[MyDeployment]
response = handle.get_user.remote(123)  # Plugin would infer DeploymentResponse[str]
user = response.result()                 # Would be typed as str

The test file includes commented-out tests that should pass once the plugin is implemented.

Tests

From master

❯ mypy python/ray/serve/tests/typing_files/check_handle_typing.py python/ray/serve/handle.py  --follow-imports=skip --ignore-missing-imports
python/ray/serve/handle.py:10: error: Module "ray._raylet" has no attribute "ObjectRefGenerator"  [attr-defined]
python/ray/serve/handle.py:162: error: Item "None" of "Optional[Any]" has no attribute "_run_router_in_separate_loop"  [union-attr]
python/ray/serve/handle.py:210: error: Item "None" of "Optional[Any]" has no attribute "assign_request"  [union-attr]
python/ray/serve/handle.py:229: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
python/ray/serve/handle.py:292: error: Unexpected keyword argument "timeout" for "result" of "Future"  [call-arg]
/home/ubuntu/.local/lib/python3.9/site-packages/mypy/typeshed/stdlib/_asyncio.pyi: note: "result" of "Future" defined here
python/ray/serve/handle.py:319: error: Incompatible types in "await" (actual type "Union[concurrent.futures._base.Future[Any], _asyncio.Future[Any]]", expected type "Awaitable[Any]")  [misc]
python/ray/serve/handle.py:803: error: Incompatible types in assignment (expression has type "type[DeploymentResponse]", variable has type "type[DeploymentResponseGenerator]")  [assignment]
python/ray/serve/tests/typing_files/check_handle_typing.py:24: error: "DeploymentResponse" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:24: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:28: error: Expression is of type "Any", not "str"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:37: error: "DeploymentResponse" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:37: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:41: error: Expression is of type "Any", not "str"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:46: error: "DeploymentResponseGenerator" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:46: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:49: error: Expression is of type "Iterator[Any]", not "Iterator[int]"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:53: error: Expression is of type "Any", not "int"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:57: error: Expression is of type "Any", not "int"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:63: error: "DeploymentResponseGenerator" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:63: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:66: error: Expression is of type "AsyncIterator[Any]", not "AsyncIterator[int]"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:70: error: Expression is of type "Any", not "int"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:74: error: Expression is of type "Any", not "int"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:88: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:88: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:97: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:104: error: "DeploymentResponse" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:104: error: "DeploymentResponseGenerator" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:115: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:115: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:119: error: Expression is of type "Any", not "DeploymentHandle"  [assert-type]
python/ray/serve/tests/typing_files/check_handle_typing.py:119: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:151: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:151: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:177: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:177: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:200: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:200: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:230: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:230: note: Error code "type-arg" not covered by "type: ignore" comment
python/ray/serve/tests/typing_files/check_handle_typing.py:262: error: "DeploymentHandle" expects no type arguments, but 1 given  [type-arg]
python/ray/serve/tests/typing_files/check_handle_typing.py:262: note: Error code "type-arg" not covered by "type: ignore" comment
Found 30 errors in 2 files (checked 2 source files)

Because of changes in this PR

❯ mypy python/ray/serve/tests/typing_files/check_handle_typing.py python/ray/serve/handle.py  --follow-imports=skip --ignore-missing-imports
python/ray/serve/handle.py:261: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
Success: no issues found in 2 source files