Add type annotation in _cuda_typerules.py
-from typing import Dict
-
import ast
+from typing import Any, Callable, Mapping, Tuple, Type
import numpy
+import numpy.typing as npt
import operator
import cupy
@@ -48,7 +48,7 @@ _scalar_gt = _core.create_comparison('scalar_less', '>')
_scalar_gte = _core.create_comparison('scalar_less', '>=')
-_py_ops = {
+_py_ops: Mapping[Type[ast.AST], Callable[..., Any]] = {
ast.And: lambda x, y: x and y,
ast.Or: lambda x, y: x or y,
ast.Add: operator.add,
@@ -75,7 +75,7 @@ _py_ops = {
}
-_numpy_ops: Dict[type, cupy.ufunc] = {
+_numpy_ops: Mapping[Type[ast.AST], cupy.ufunc] = {
ast.And: ops.logical_and,
ast.Or: ops.logical_or,
ast.Add: arithmetic.add,
@@ -102,11 +102,11 @@ _numpy_ops: Dict[type, cupy.ufunc] = {
}
-def get_pyfunc(op_type):
+def get_pyfunc(op_type: Type[ast.AST]) -> Callable[..., Any]:
return _py_ops[op_type]
-def get_ufunc(mode, op_type):
+def get_ufunc(mode: str, op_type: Type[ast.AST]) -> cupy.ufunc:
if mode == 'numpy':
return _numpy_ops[op_type]
if mode == 'cuda':
@@ -114,7 +114,7 @@ def get_ufunc(mode, op_type):
assert False
-def get_ctype_from_scalar(mode, x):
+def get_ctype_from_scalar(mode: str, x: Any) -> _cuda_types.Scalar:
if isinstance(x, numpy.generic):
return _cuda_types.Scalar(x.dtype)
@@ -147,13 +147,18 @@ def get_ctype_from_scalar(mode, x):
_typechars = '?bBhHiIlLqQefdFD'
-def _cuda_can_cast(from_dtype, to_dtype):
+def _cuda_can_cast(from_dtype: npt.DTypeLike, to_dtype: npt.DTypeLike) -> bool:
from_dtype = numpy.dtype(from_dtype)
to_dtype = numpy.dtype(to_dtype)
return _typechars.find(from_dtype.char) <= _typechars.find(to_dtype.char)
-def guess_routine(ufunc, in_types, dtype, mode):
+def guess_routine(
+ ufunc: cupy.ufunc,
+ in_types: Tuple[numpy.dtype],
+ dtype: numpy.dtype,
+ mode: str,
+) -> cupy._core._kernel._Op:
if dtype is not None:
return ufunc._ops._guess_routine_from_dtype(dtype)
can_cast = numpy.can_cast if mode == 'numpy' else _cuda_can_cast
