From bb4472eccbd3bcb76db362dcfb615b3c91ecc275 Mon Sep 17 00:00:00 2001 From: Nicolas Kruse Date: Mon, 1 Dec 2025 00:12:22 +0100 Subject: [PATCH] variable type hints updated --- src/copapy/_basic_types.py | 48 +++++++++++++++++++++++++++++--------- 1 file changed, 37 insertions(+), 11 deletions(-) diff --git a/src/copapy/_basic_types.py b/src/copapy/_basic_types.py index 1d8ce67..5170d04 100644 --- a/src/copapy/_basic_types.py +++ b/src/copapy/_basic_types.py @@ -104,6 +104,8 @@ class variable(Generic[TNum], Net): self.source = CPConstant(source) self.dtype = 'int' + @overload + def __add__(self: 'variable[TNum]', other: 'variable[TNum] | TNum') -> 'variable[TNum]': ... @overload def __add__(self: 'variable[int]', other: uniint) -> 'variable[int]': ... @overload @@ -117,6 +119,8 @@ class variable(Generic[TNum], Net): return self return add_op('add', [self, other], True) + @overload + def __radd__(self: 'variable[TNum]', other: TNum) -> 'variable[TNum]': ... @overload def __radd__(self: 'variable[int]', other: int) -> 'variable[int]': ... @overload @@ -126,6 +130,8 @@ class variable(Generic[TNum], Net): return self return add_op('add', [self, other], True) + @overload + def __sub__(self: 'variable[TNum]', other: 'variable[TNum] | TNum') -> 'variable[TNum]': ... @overload def __sub__(self: 'variable[int]', other: uniint) -> 'variable[int]': ... @overload @@ -137,6 +143,8 @@ class variable(Generic[TNum], Net): def __sub__(self, other: TVarNumb) -> Any: return add_op('sub', [self, other]) + @overload + def __rsub__(self: 'variable[TNum]', other: TNum) -> 'variable[TNum]': ... @overload def __rsub__(self: 'variable[int]', other: int) -> 'variable[int]': ... @overload @@ -144,6 +152,8 @@ class variable(Generic[TNum], Net): def __rsub__(self, other: NumLike) -> Any: return add_op('sub', [other, self]) + @overload + def __mul__(self: 'variable[TNum]', other: 'variable[TNum] | TNum') -> 'variable[TNum]': ... @overload def __mul__(self: 'variable[int]', other: uniint) -> 'variable[int]': ... @overload @@ -155,6 +165,8 @@ class variable(Generic[TNum], Net): def __mul__(self, other: TVarNumb) -> Any: return add_op('mul', [self, other], True) + @overload + def __rmul__(self: 'variable[TNum]', other: TNum) -> 'variable[TNum]': ... @overload def __rmul__(self: 'variable[int]', other: int) -> 'variable[int]': ... @overload @@ -168,6 +180,8 @@ class variable(Generic[TNum], Net): def __rtruediv__(self, other: NumLike) -> 'variable[float]': return add_op('div', [other, self]) + @overload + def __floordiv__(self: 'variable[TNum]', other: 'variable[TNum] | TNum') -> 'variable[TNum]': ... @overload def __floordiv__(self: 'variable[int]', other: uniint) -> 'variable[int]': ... @overload @@ -179,6 +193,8 @@ class variable(Generic[TNum], Net): def __floordiv__(self, other: TVarNumb) -> Any: return add_op('floordiv', [self, other]) + @overload + def __rfloordiv__(self: 'variable[TNum]', other: TNum) -> 'variable[TNum]': ... @overload def __rfloordiv__(self: 'variable[int]', other: int) -> 'variable[int]': ... @overload @@ -187,7 +203,9 @@ class variable(Generic[TNum], Net): return add_op('floordiv', [other, self]) def __neg__(self: TCPNum) -> TCPNum: - return cast(TCPNum, add_op('sub', [variable(0), self])) + if self.dtype == 'int': + return cast(TCPNum, add_op('sub', [variable(0), self])) + return cast(TCPNum, add_op('sub', [variable(0.0), self])) def __gt__(self, other: TVarNumb) -> 'variable[int]': ret = add_op('gt', [self, other]) @@ -214,36 +232,44 @@ class variable(Generic[TNum], Net): return variable(ret.source, dtype='bool') @overload - def __mod__(self, other: TCPNum) -> TCPNum: ... + def __mod__(self: 'variable[TNum]', other: 'variable[TNum] | TNum') -> 'variable[TNum]': ... @overload - def __mod__(self: TCPNum, other: uniint) -> TCPNum: ... + def __mod__(self: 'variable[int]', other: uniint) -> 'variable[int]': ... @overload def __mod__(self, other: unifloat) -> 'variable[float]': ... @overload - def __mod__(self, other: NumLike) -> 'variable[float] | variable[int]': ... - def __mod__(self, other: NumLike) -> Any: + def __mod__(self: 'variable[float]', other: NumLike) -> 'variable[float]': ... + @overload + def __mod__(self, other: TVarNumb) -> 'variable[float] | variable[int]': ... + def __mod__(self, other: TVarNumb) -> Any: return add_op('mod', [self, other]) @overload - def __rmod__(self: TCPNum, other: int) -> TCPNum: ... + def __rmod__(self: 'variable[TNum]', other: TNum) -> 'variable[TNum]': ... + @overload + def __rmod__(self: 'variable[int]', other: int) -> 'variable[int]': ... @overload def __rmod__(self, other: float) -> 'variable[float]': ... def __rmod__(self, other: NumLike) -> Any: return add_op('mod', [other, self]) @overload - def __pow__(self, other: TCPNum) -> TCPNum: ... + def __pow__(self: 'variable[TNum]', other: 'variable[TNum] | TNum') -> 'variable[TNum]': ... @overload - def __pow__(self: TCPNum, other: uniint) -> TCPNum: ... + def __pow__(self: 'variable[int]', other: uniint) -> 'variable[int]': ... @overload def __pow__(self, other: unifloat) -> 'variable[float]': ... @overload - def __pow__(self, other: NumLike) -> 'variable[float] | variable[int]': ... - def __pow__(self, other: NumLike) -> Any: + def __pow__(self: 'variable[float]', other: NumLike) -> 'variable[float]': ... + @overload + def __pow__(self, other: TVarNumb) -> 'variable[float] | variable[int]': ... + def __pow__(self, other: TVarNumb) -> Any: return cp.pow(self, other) @overload - def __rpow__(self: TCPNum, other: int) -> TCPNum: ... + def __rpow__(self: 'variable[TNum]', other: TNum) -> 'variable[TNum]': ... + @overload + def __rpow__(self: 'variable[int]', other: int) -> 'variable[int]': ... @overload def __rpow__(self, other: float) -> 'variable[float]': ... def __rpow__(self, other: NumLike) -> Any: