fix(pyargus): address typing issues

This addresses some of the issues with inheritance (internal to the rust
module) for signals, and generally making mypy and flake8 happy.

pyargus/argus/__init__.py

@@ -5,7 +5,7 @@ from typing import List, Optional, Tuple, Type, Union
 from argus import _argus
 from argus._argus import DType as DType
 from argus.exprs import ConstBool, ConstFloat, ConstInt, ConstUInt, VarBool, VarFloat, VarInt, VarUInt
-from argus.signals import BoolSignal, FloatSignal, IntSignal, Signal, UnsignedIntSignal
+from argus.signals import BoolSignal, FloatSignal, IntSignal, UnsignedIntSignal
 
 try:
     __doc__ = _argus.__doc__
@@ -51,26 +51,23 @@ def signal(
     dtype: Union[DType, Type[AllowedDtype]],
     *,
     data: Optional[Union[AllowedDtype, List[Tuple[float, AllowedDtype]]]] = None,
-) -> Signal:
+) -> Union[BoolSignal, UnsignedIntSignal, IntSignal, FloatSignal]:
     """Create a signal of the given type
 
     Parameters
     ----------
 
+    dtype:
+        Type of the signal
+
     data :
         If a constant scalar is given, a constant signal is created. Otherwise, if a list of sample points are given, a sampled
         signal is constructed. Otherwise, an empty signal is created.
     """
-    if isinstance(dtype, type):
-        if dtype == bool:
-            dtype = DType.Bool
-        elif dtype == int:
-            dtype = DType.Int
-        elif dtype == float:
-            dtype = DType.Float
-
     factory: Type[Union[BoolSignal, UnsignedIntSignal, IntSignal, FloatSignal]]
     expected_type: Type[AllowedDtype]
+
+    dtype = DType.convert(dtype)
     if dtype == DType.Bool:
         factory = BoolSignal
         expected_type = bool
@@ -89,10 +86,9 @@ def signal(
     if data is None:
         return factory.from_samples([])
     elif isinstance(data, (list, tuple)):
-        return factory.from_samples(data)  # type: ignore
-    else:
-        assert isinstance(data, expected_type)
-        return factory.constant(data)  # type: ignore
+        return factory.from_samples(data)  # type: ignore[arg-type]
+    assert isinstance(data, expected_type)
+    return factory.constant(data)  # type: ignore[arg-type]
 
 
 __all__ = [

pyargus/argus/_argus.pyi

@@ -1,8 +1,8 @@
-from abc import ABC
-from enum import Enum, auto
-from typing import List, Optional, Tuple, final
+from typing import ClassVar, Protocol, TypeVar, final
 
-class NumExpr(ABC):
+from typing_extensions import Self
+
+class NumExpr(Protocol):
     def __ge__(self, other) -> NumExpr: ...
     def __gt__(self, other) -> NumExpr: ...
     def __le__(self, other) -> NumExpr: ...
@@ -50,11 +50,11 @@ class Negate(NumExpr):
 
 @final
 class Add(NumExpr):
-    def __init__(self, args: List[NumExpr]): ...
+    def __init__(self, args: list[NumExpr]): ...
 
 @final
 class Mul(NumExpr):
-    def __init__(self, args: List[NumExpr]): ...
+    def __init__(self, args: list[NumExpr]): ...
 
 @final
 class Div(NumExpr):
@@ -64,7 +64,7 @@ class Div(NumExpr):
 class Abs(NumExpr):
     def __init__(self, arg: NumExpr): ...
 
-class BoolExpr(ABC):
+class BoolExpr(Protocol):
     def __and__(self, other) -> BoolExpr: ...
     def __invert__(self) -> BoolExpr: ...
     def __or__(self, other) -> BoolExpr: ...
@@ -100,11 +100,11 @@ class Not(BoolExpr):
 
 @final
 class And(BoolExpr):
-    def __init__(self, args: List[BoolExpr]): ...
+    def __init__(self, args: list[BoolExpr]): ...
 
 @final
 class Or(BoolExpr):
-    def __init__(self, args: List[BoolExpr]): ...
+    def __init__(self, args: list[BoolExpr]): ...
 
 @final
 class Next(BoolExpr):
@@ -123,57 +123,63 @@ class Until(BoolExpr):
     def __init__(self, lhs: BoolExpr, rhs: BoolExpr): ...
 
 @final
-class DType(Enum):
-    Bool = auto()
-    Int = auto()
-    UnsignedInt = auto()
-    Float = auto()
+class DType:
+    Bool: ClassVar[DType] = ...
+    Float: ClassVar[DType] = ...
+    Int: ClassVar[DType] = ...
+    UnsignedInt: ClassVar[DType] = ...
 
-class Signal(ABC): ...
+    @classmethod
+    def convert(cls, dtype: type[bool | int | float] | Self) -> Self: ...  # noqa: Y041
+    def __eq__(self, other) -> bool: ...
+    def __int__(self) -> int: ...
+
+_SignalKind = TypeVar("_SignalKind", bool, int, float, covariant=True)
+
+class Signal(Protocol[_SignalKind]):
+    def is_empty(self) -> bool: ...
+    @property
+    def start_time(self) -> float | None: ...
+    @property
+    def end_time(self) -> float | None: ...
+    @property
+    def kind(self) -> type[bool | int | float]: ...
 
 @final
-class BoolSignal(Signal):
-    def __init__(self): ...
-    @staticmethod
-    def constant(value: bool) -> BoolSignal: ...
-    @staticmethod
-    def from_samples(samples: List[Tuple[float, bool]]) -> BoolSignal: ...
+class BoolSignal(Signal[bool]):
+    @classmethod
+    def constant(cls, value: bool) -> Self: ...
+    @classmethod
+    def from_samples(cls, samples: list[tuple[float, bool]]) -> Self: ...
     def push(self, time: float, value: bool): ...
-    def is_empty(self) -> bool: ...
-    def at(self, time: float) -> Optional[bool]: ...
+    def at(self, time: float) -> _SignalKind | None: ...
 
 @final
-class IntSignal(Signal):
-    def __init__(self): ...
-    @staticmethod
-    def constant(value: int) -> IntSignal: ...
-    @staticmethod
-    def from_samples(samples: List[Tuple[float, int]]) -> IntSignal: ...
+class IntSignal(Signal[int]):
+    @classmethod
+    def constant(cls, value: int) -> Self: ...
+    @classmethod
+    def from_samples(cls, samples: list[tuple[float, int]]) -> Self: ...
     def push(self, time: float, value: int): ...
-    def is_empty(self) -> bool: ...
-    def at(self, time: float) -> Optional[int]: ...
+    def at(self, time: float) -> int | None: ...
 
 @final
-class UnsignedIntSignal(Signal):
-    def __init__(self): ...
-    @staticmethod
-    def constant(value: int) -> UnsignedIntSignal: ...
-    @staticmethod
-    def from_samples(samples: List[Tuple[float, int]]) -> UnsignedIntSignal: ...
+class UnsignedIntSignal(Signal[int]):
+    @classmethod
+    def constant(cls, value: int) -> Self: ...
+    @classmethod
+    def from_samples(cls, samples: list[tuple[float, int]]) -> Self: ...
     def push(self, time: float, value: int): ...
-    def is_empty(self) -> bool: ...
-    def at(self, time: float) -> Optional[int]: ...
+    def at(self, time: float) -> int | None: ...
 
 @final
-class FloatSignal(Signal):
-    def __init__(self): ...
-    @staticmethod
-    def constant(value: float) -> FloatSignal: ...
-    @staticmethod
-    def from_samples(samples: List[Tuple[float, float]]) -> FloatSignal: ...
+class FloatSignal(Signal[float]):
+    @classmethod
+    def constant(cls, value: float) -> Self: ...
+    @classmethod
+    def from_samples(cls, samples: list[tuple[float, float]]) -> Self: ...
     def push(self, time: float, value: float): ...
-    def is_empty(self) -> bool: ...
-    def at(self, time: float) -> Optional[float]: ...
+    def at(self, time: float) -> float | None: ...
 
 @final
 class Trace: ...