variable type hints updated

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: