[Core] Type check ObjectRef (#9856)

* Type check ObjectRef

* Bug fix

* Port typing tests to bazel test
This commit is contained in:
Simon Mo
2020-08-11 10:38:29 -07:00
committed by GitHub
parent 98df612010
commit 4c52adddfa
8 changed files with 119 additions and 63 deletions
+37 -28
View File
@@ -1,4 +1,5 @@
from typing import Any, Callable, Generic, Optional, TypeVar, Union, overload
# yapf: disable
from typing import Any, Callable, Generic, Optional, TypeVar, Union, overload, Sequence, List
from ray._raylet import ObjectRef
@@ -16,58 +17,66 @@ T9 = TypeVar("T9")
R = TypeVar("R")
class RemoteFunction(Generic[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]):
def __init__(self, function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9], Any]) -> None: pass
class RemoteFunction(Generic[R, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]):
def __init__(self, function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9], R]) -> None: pass
@overload
def remote(self) -> ObjectRef: ...
def remote(self) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef], arg3: Union[T3, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]], arg3: Union[T3, ObjectRef[T3]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef], arg3: Union[T3, ObjectRef], arg4: Union[T4, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]], arg3: Union[T3, ObjectRef[T3]], arg4: Union[T4, ObjectRef[T4]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef], arg3: Union[T3, ObjectRef], arg4: Union[T4, ObjectRef], arg5: Union[T5, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]], arg3: Union[T3, ObjectRef[T3]], arg4: Union[T4, ObjectRef[T4]], arg5: Union[T5, ObjectRef[T5]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef], arg3: Union[T3, ObjectRef], arg4: Union[T4, ObjectRef], arg5: Union[T5, ObjectRef], arg6: Union[T6, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]], arg3: Union[T3, ObjectRef[T3]], arg4: Union[T4, ObjectRef[T4]], arg5: Union[T5, ObjectRef[T5]], arg6: Union[T6, ObjectRef[T6]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef], arg3: Union[T3, ObjectRef], arg4: Union[T4, ObjectRef], arg5: Union[T5, ObjectRef], arg6: Union[T6, ObjectRef], arg7: Union[T7, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]], arg3: Union[T3, ObjectRef[T3]], arg4: Union[T4, ObjectRef[T4]], arg5: Union[T5, ObjectRef[T5]], arg6: Union[T6, ObjectRef[T6]], arg7: Union[T7, ObjectRef[T7]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef], arg3: Union[T3, ObjectRef], arg4: Union[T4, ObjectRef], arg5: Union[T5, ObjectRef], arg6: Union[T6, ObjectRef], arg7: Union[T7, ObjectRef], arg8: Union[T8, ObjectRef]) -> ObjectRef: ...
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]], arg3: Union[T3, ObjectRef[T3]], arg4: Union[T4, ObjectRef[T4]], arg5: Union[T5, ObjectRef[T5]], arg6: Union[T6, ObjectRef[T6]], arg7: Union[T7, ObjectRef[T7]], arg8: Union[T8, ObjectRef[T8]]) -> ObjectRef[R]: ...
@overload
def remote(self, arg0: Union[T0, ObjectRef], arg1: Union[T1, ObjectRef], arg2: Union[T2, ObjectRef], arg3: Union[T3, ObjectRef], arg4: Union[T4, ObjectRef], arg5: Union[T5, ObjectRef], arg6: Union[T6, ObjectRef], arg7: Union[T7, ObjectRef], arg8: Union[T8, ObjectRef], arg9: Union[T9, ObjectRef]) -> ObjectRef: ...
def remote(self, *args, **kwargs) -> ObjectRef:
pass
def remote(self, arg0: Union[T0, ObjectRef[T0]], arg1: Union[T1, ObjectRef[T1]], arg2: Union[T2, ObjectRef[T2]], arg3: Union[T3, ObjectRef[T3]], arg4: Union[T4, ObjectRef[T4]], arg5: Union[T5, ObjectRef[T5]], arg6: Union[T6, ObjectRef[T6]], arg7: Union[T7, ObjectRef[T7]], arg8: Union[T8, ObjectRef[T8]], arg9: Union[T9, ObjectRef[T9]]) -> ObjectRef[R]: ...
def remote(self, *args, **kwargs) -> ObjectRef[R]: ...
@overload
def remote(function: Callable[[], R]) -> RemoteFunction[None, None, None, None, None, None, None, None, None, None]: ...
def remote(function: Callable[[], R]) -> RemoteFunction[R, None, None, None, None, None, None, None, None, None, None]: ...
@overload
def remote(function: Callable[[T0], R]) -> RemoteFunction[T0, None, None, None, None, None, None, None, None, None]: ...
def remote(function: Callable[[T0], R]) -> RemoteFunction[R, T0, None, None, None, None, None, None, None, None, None]: ...
@overload
def remote(function: Callable[[T0, T1], R]) -> RemoteFunction[T0, T1, None, None, None, None, None, None, None, None]: ...
def remote(function: Callable[[T0, T1], R]) -> RemoteFunction[R, T0, T1, None, None, None, None, None, None, None, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2], R]) -> RemoteFunction[T0, T1, T2, None, None, None, None, None, None, None]: ...
def remote(function: Callable[[T0, T1, T2], R]) -> RemoteFunction[R, T0, T1, T2, None, None, None, None, None, None, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2, T3], R]) -> RemoteFunction[T0, T1, T2, T3, None, None, None, None, None, None]: ...
def remote(function: Callable[[T0, T1, T2, T3], R]) -> RemoteFunction[R, T0, T1, T2, T3, None, None, None, None, None, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2, T3, T4], R]) -> RemoteFunction[T0, T1, T2, T3, T4, None, None, None, None, None]: ...
def remote(function: Callable[[T0, T1, T2, T3, T4], R]) -> RemoteFunction[R, T0, T1, T2, T3, T4, None, None, None, None, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2, T3, T4, T5], R]) -> RemoteFunction[T0, T1, T2, T3, T4, T5, None, None, None, None]: ...
def remote(function: Callable[[T0, T1, T2, T3, T4, T5], R]) -> RemoteFunction[R, T0, T1, T2, T3, T4, T5, None, None, None, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6], R]) -> RemoteFunction[T0, T1, T2, T3, T4, T5, T6, None, None, None]: ...
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6], R]) -> RemoteFunction[R, T0, T1, T2, T3, T4, T5, T6, None, None, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7], R]) -> RemoteFunction[T0, T1, T2, T3, T4, T5, T6, T7, None, None]: ...
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7], R]) -> RemoteFunction[R, T0, T1, T2, T3, T4, T5, T6, T7, None, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8], R]) -> RemoteFunction[T0, T1, T2, T3, T4, T5, T6, T7, T8, None]: ...
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8], R]) -> RemoteFunction[R, T0, T1, T2, T3, T4, T5, T6, T7, T8, None]: ...
@overload
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9], R]) -> RemoteFunction[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]: ...
def remote(function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9], R]) -> RemoteFunction[R, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]: ...
# Pass on typing actors for now. The following makes it so no type errors are generated for actors.
@overload
def remote(t: type) -> Any: ...
def remote(function: Callable[..., R]) -> RemoteFunction[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]: pass
def remote(function: Callable[..., R]) -> RemoteFunction[R, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]: ...
@overload
def get(object_refs: Sequence[ObjectRef[Any]], timeout: Optional[float] = None) -> List[Any]: ...
@overload
def get(object_refs: Sequence[ObjectRef[R]], timeout: Optional[float] = None) -> List[R]: ...
@overload
def get(object_refs: ObjectRef[R], timeout: Optional[float] = None) -> R: ...