vector typing advanced

src/copapy/__init__.py

@@ -1,13 +1,14 @@
 from ._target import Target
 from ._basic_types import NumLike, variable, \
-    CPNumber, cpvector, generic_sdb, iif
+    CPNumber, generic_sdb, iif
+from ._vectors import vector
 
 __all__ = [
     "Target",
     "NumLike",
     "variable",
     "CPNumber",
-    "cpvector",
     "generic_sdb",
     "iif",
+    "vector"
 ]

src/copapy/_vectors.py

@@ -1,14 +1,45 @@
-from copapy import NumLike, CPNumber, cpint, cpfloat, cpbool
+from numpy import isin
-from typing import Generic, TypeVar, Iterable, Any
+from copapy import NumLike, CPNumber, variable
+from typing import Generic, TypeVar, Iterable, Any, overload
 
-T = TypeVar("T", bound=CPNumber)
+from copapy._basic_types import TNum
+
+T = TypeVar("T", int, float, bool)
 T2 = TypeVar("T2", bound=CPNumber)
 
-class cpvector(Generic[T]):
+class vector(Generic[T]):
-    def __init__(self, value: Iterable[T]):
+    def __init__(self, values: Iterable[T | variable[T]]):
-        self.value = tuple(value)
+        #self.values: tuple[variable[T], ...] = tuple(v if isinstance(v, variable) else variable(v) for v in values)
+        self.values: tuple[variable[T] | T, ...] = tuple(values)
+
+    @overload
+    def __add__(self, other: 'vector[float] | variable[float] | float') -> 'vector[float]':
+        ...
+
+    @overload
+    def __add__(self: 'vector[T]', other: 'vector[int] | variable[int] | int') -> 'vector[T]':
+        ...
+
+    def __add__(self, other: 'vector[Any] | variable[Any] | float | int') -> Any:
+        if isinstance(other, vector):
+            assert len(self.values) == len(other.values)
+            return vector(a + b for a, b in zip(self.values, other.values))
+        else:
+            return vector(a + other for a in self.values)
+    
+    #@overload
+    #def sum(self: 'vector[float]') -> variable[float]:
+    #    ...
+
+    #@overload
+    #def sum(self: 'vector[int]') -> variable[int]:
+    #    ...
+
+    #def sum(self: 'vector[T]') -> variable[T] | T:
+    #    comp_time = sum(v for v in self.values if not isinstance(v, variable))
+    #    run_time = sum(v for v in self.values if isinstance(v, variable))
+    #    if isinstance(run_time, variable):
+    #        return comp_time + run_time  # type: ignore
+    #    else:
+    #        return comp_time
     
-    def __add__(self, other: 'cpvector[Any]') -> 'cpvector[CPNumber]':
-        assert len(self.value) == len(other.value)
-        tup = (a + b for a, b in zip(self.value, other.value))
-        return cpvector(*(v for v in tup if isinstance(v, CPNumber)))

tests/test_vector.py

@@ -0,0 +1,5 @@
+import copapy as cp
+
+def test_vec():
+    tt = cp.vector(range(3)) + cp.vector([1.1,2.2,3.3])
+    tt2 = (cp.vector(range(3)) + 5.6)