Add type annotation in _internal_types.py
import itertools
+from typing import Any, NoReturn, Optional, Union, TYPE_CHECKING
from cupyx.jit import _cuda_types
from cupyx.jit import _cuda_typerules
+if TYPE_CHECKING:
+ from cupyx.jit._compile import Environment
+
class Expr:
- pass
+
+ def __repr__(self) -> str:
+ raise NotImplementedError
class Data(Expr):
- def __init__(self, code: str, ctype: _cuda_types.TypeBase):
+
+ def __init__(self, code: str, ctype: _cuda_types.TypeBase) -> None:
assert isinstance(code, str)
assert isinstance(ctype, _cuda_types.TypeBase)
self.code = code
@@ -23,11 +30,11 @@ class Data(Expr):
def obj(self):
raise ValueError(f'Constant value is requried: {self.code}')
- def __repr__(self):
+ def __repr__(self) -> str:
return f'<Data code = "{self.code}", type = {self.ctype}>'
@classmethod
- def init(cls, x: Expr, env):
+ def init(cls, x: Expr, env) -> 'Data':
if isinstance(x, Data):
return x
if isinstance(x, Constant):
@@ -38,21 +45,27 @@ class Data(Expr):
class Constant(Expr):
- def __init__(self, obj):
+
+ def __init__(self, obj: Any) -> None:
self._obj = obj
@property
- def obj(self):
+ def obj(self) -> Any:
return self._obj
- def __repr__(self):
+ def __repr__(self) -> str:
return f'<Constant obj = "{self.obj}">'
class Range(Expr):
- def __init__(self, start, stop, step, ctype, step_is_positive, *,
- unroll=None):
+ def __init__(
+ self, start: Expr, stop: Expr, step: Expr,
+ ctype: _cuda_types.Scalar,
+ step_is_positive: Optional[bool],
+ *,
+ unroll: Union[None, int, bool] = None,
+ ) -> None:
self.start = start
self.stop = stop
self.step = step
@@ -66,21 +79,21 @@ class BuiltinFunc(Expr):
# - either call or call_const
# - `__call__` with a correct signature, which calls the parent's __call__
- def call(self, env, *args, **kwargs):
+ def call(self, env: 'Environment', *args, **kwargs) -> Expr:
for x in itertools.chain(args, kwargs.values()):
if not isinstance(x, Constant):
raise TypeError('Arguments must be constants.')
- args = [x.obj for x in args]
+ args = tuple([x.obj for x in args])
kwargs = dict([(k, v.obj) for k, v in kwargs.items()])
return self.call_const(env, *args, **kwargs)
- def call_const(self, env, *args, **kwarg):
+ def call_const(self, env: 'Environment', *args: Any, **kwarg: Any) -> Expr:
raise NotImplementedError
- def __init__(self):
+ def __init__(self) -> None:
self.__doc__ = type(self).__call__.__doc__
- def __call__(self):
+ def __call__(self) -> NoReturn:
raise RuntimeError('Cannot call this function from Python layer.')
@classmethod
