chore: import upstream snapshot with attribution

python/ray/tests/typing_files/check_typing_actor_async.py

@@ -0,0 +1,37 @@
+import ray
+from ray import ObjectRef
+
+
+@ray.remote
+class AsyncActor:
+    @ray.method
+    async def add(self, a: int, b: int) -> int:
+        return a + b
+
+    @ray.method(num_returns=1)
+    async def mul(self, a: int, b: int) -> int:
+        return a * b
+
+    @ray.method(num_returns=1)
+    def divide(self, a: int, b: int) -> int:
+        if b == 0:
+            raise ValueError("Division by zero")
+        return a // b
+
+    @ray.method(num_returns=1)
+    def echo(self, x: str) -> str:
+        return x
+
+
+actor = AsyncActor.remote()
+
+ref_add: ObjectRef[int] = actor.add.remote(1, 2)
+ref_mul: ObjectRef[int] = actor.mul.remote(2, 3)
+ref_echo: ObjectRef[str] = actor.echo.remote("hello")
+ref_divide: ObjectRef[int] = actor.divide.remote(10, 2)
+
+# ray.get() should resolve to int for both
+result_add: int = ray.get(ref_add)
+result_mul: int = ray.get(ref_mul)
+result_echo: str = ray.get(ref_echo)
+result_divide: int = ray.get(ref_divide)

python/ray/tests/typing_files/check_typing_bad.py

@@ -0,0 +1,33 @@
+import ray
+
+ray.init()
+
+
+@ray.remote
+def f(a: int) -> str:
+    return "a = {}".format(a + 1)
+
+
+@ray.remote
+def g(s: str) -> str:
+    return s + " world"
+
+
+@ray.remote
+def h(a: str, b: int) -> str:
+    return a
+
+
+# Does not typecheck due to incorrect input type:
+a = h.remote(1, 1)
+b = f.remote("hello")
+c = f.remote(1, 1)
+d = f.remote(1) + 1
+
+# Check return type
+ref_to_str = f.remote(1)
+unwrapped_str = ray.get(ref_to_str)
+unwrapped_str + 100  # Fail
+
+# Check ObjectRef[T] as args
+f.remote(ref_to_str)  # Fail

python/ray/tests/typing_files/check_typing_good.py

@@ -0,0 +1,116 @@
+from typing import Generator
+
+import ray
+from ray import ObjectRef
+
+ray.init()
+
+with ray.init(num_cpus=1, num_gpus=0):
+    pass
+
+
+@ray.remote
+def int_task() -> int:
+    return 1
+
+
+@ray.remote
+def f(a: int) -> str:
+    return "a = {}".format(a + 1)
+
+
+@ray.remote
+def g(s: str) -> str:
+    return s + " world"
+
+
+@ray.remote
+def h(a: str, b: int) -> str:
+    return a
+
+
+def func(a: "ObjectRef[str]"):
+    pass
+
+
+# Make sure the function arg is check
+print(f.remote(1))
+object_ref_str = f.remote(1)
+object_ref_int = int_task.remote()
+
+# Make sure the ObjectRef[T] variant of function arg is checked
+print(g.remote(object_ref_str))
+
+# Make sure it is backward compatible after
+# introducing generator types.
+func(object_ref_str)
+
+# Make sure there can be mixed T0 and ObjectRef[T1] for args
+print(h.remote(object_ref_str, 100))
+
+ready, unready = ray.wait([object_ref_str, object_ref_int])
+
+# Make sure the return type is checked.
+xy = ray.get(object_ref_str) + "y"
+
+
+# Right now, we only check if it doesn't raise errors.
+@ray.remote
+def generator_1() -> Generator[int, None, None]:
+    yield 1
+
+
+gen = generator_1.remote()
+
+
+"""
+TODO(sang): Enable it.
+Test generator
+
+Generator can have 4 different output
+per generator and async generator. See
+https://docs.python.org/3/library/typing.html#typing.Generator
+for more details.
+"""
+
+# @ray.remote
+# def generator_1() -> Generator[int, None, None]:
+#     yield 1
+
+
+# @ray.remote
+# def generator_2() -> Iterator[int]:
+#     yield 1
+
+
+# @ray.remote
+# def generator_3() -> Iterable[int]:
+#     yield 1
+
+
+# gen: ObjectRefGenerator[int] = generator_1.remote()
+# gen2: ObjectRefGenerator[int] = generator_2.remote()
+# gen3: ObjectRefGenerator[int] = generator_3.remote()
+
+
+# next_item: ObjectRef[int] = gen.__next__()
+
+
+# @ray.remote
+# async def async_generator_1() -> AsyncGenerator[int, None]:
+#     yield 1
+
+
+# @ray.remote
+# async def async_generator_2() -> AsyncIterator[int]:
+#     yield 1
+
+
+# @ray.remote
+# async def async_generator_3() -> AsyncIterable[int]:
+#     yield 1
+
+
+# gen4: ObjectRefGenerator[int] = async_generator_1.remote()
+# gen5: ObjectRefGenerator[int] = async_generator_2.remote()
+# gen6: ObjectRefGenerator[int] = async_generator_3.remote()