refactor(py): update the type stubs (and expose DType)

justfile

@@ -6,6 +6,7 @@ test *ARGS:
 
 check:
   cargo +nightly clippy
+  cd pyargus && stubtest argus
   cd pyargus && mypy .
   cd pyargus && flake8
   cd pyargus && ruff .

pyargus/argus/__init__.py

@@ -1,8 +1,9 @@
 from argus import _argus
-from argus._argus import *  # noqa: F401
-
-__all__ = []
 
+try:
     __doc__ = _argus.__doc__
-if hasattr(_argus, "__all__"):
+except AttributeError:
-    __all__ += _argus.__all__
+    ...
+
+
+

pyargus/argus/__init__.pyi

@@ -1,30 +0,0 @@
-from argus._argus import *
-
-# Names in __all__ with no definition:
-#   Add
-#   Always
-#   And
-#   BoolExpr
-#   BoolSignal
-#   Cmp
-#   ConstBool
-#   ConstFloat
-#   ConstInt
-#   ConstUInt
-#   Div
-#   Eventually
-#   FloatSignal
-#   IntSignal
-#   Mul
-#   Negate
-#   Next
-#   Not
-#   NumExpr
-#   Or
-#   Signal
-#   UnsignedIntSignal
-#   Until
-#   VarBool
-#   VarFloat
-#   VarInt
-#   VarUInt

pyargus/argus/_argus.pyi

@@ -1,4 +1,5 @@
 from abc import ABC
+from enum import Enum, auto
 from typing import List, Tuple, final
 
 class NumExpr(ABC):
@@ -121,6 +122,13 @@ class Eventually(BoolExpr):
 class Until(BoolExpr):
     def __init__(self, lhs: BoolExpr, rhs: BoolExpr): ...
 
+@final
+class DType(Enum):
+    Bool = auto()
+    Int = auto()
+    UnsignedInt = auto()
+    Float = auto()
+
 class Signal(ABC): ...
 
 @final
@@ -162,7 +170,5 @@ class FloatSignal(Signal):
 @final
 class Trace: ...
 
-@final
+def eval_bool_semantics(expr: BoolExpr, trace: Trace) -> BoolSignal: ...
-class BooleanSemantics:
+def eval_robust_semantics(expr: BoolExpr, trace: Trace) -> BoolSignal: ...
-    @staticmethod
-    def eval(expr: BoolExpr, trace: Trace) -> BoolSignal: ...

pyargus/argus/exprs.py

@@ -0,0 +1,49 @@
+from argus._argus import (
+    Abs,
+    Add,
+    Always,
+    And,
+    BoolExpr,
+    ConstBool,
+    ConstFloat,
+    ConstInt,
+    ConstUInt,
+    Div,
+    Eventually,
+    Mul,
+    Negate,
+    Next,
+    Not,
+    NumExpr,
+    Or,
+    Until,
+    VarBool,
+    VarFloat,
+    VarInt,
+    VarUInt,
+)
+
+__all__ = [
+    "Abs",
+    "Add",
+    "Always",
+    "And",
+    "BoolExpr",
+    "ConstBool",
+    "ConstFloat",
+    "ConstInt",
+    "ConstUInt",
+    "Div",
+    "Eventually",
+    "Mul",
+    "Negate",
+    "Next",
+    "Not",
+    "NumExpr",
+    "Or",
+    "Until",
+    "VarBool",
+    "VarFloat",
+    "VarInt",
+    "VarUInt",
+]

pyargus/argus/semantics.py

@@ -0,0 +1,7 @@
+from argus._argus import Trace, eval_bool_semantics, eval_robust_semantics
+
+__all__ = [
+    "Trace",
+    "eval_bool_semantics",
+    "eval_robust_semantics",
+]

pyargus/argus/signals.py

@@ -0,0 +1,9 @@
+from argus._argus import BoolSignal, FloatSignal, IntSignal, Signal, UnsignedIntSignal
+
+__all__ = [
+    "Signal",
+    "BoolSignal",
+    "IntSignal",
+    "UnsignedIntSignal",
+    "FloatSignal",
+]

pyargus/src/lib.rs

@@ -27,11 +27,22 @@ impl From<PyArgusError> for PyErr {
     }
 }
 
+#[pyclass]
+#[derive(Copy, Clone, Debug)]
+pub enum DType {
+    Bool,
+    Int,
+    UnsignedInt,
+    Float,
+}
+
 #[pymodule]
 #[pyo3(name = "_argus")]
 fn pyargus(py: Python, m: &PyModule) -> PyResult<()> {
     pyo3_log::init();
 
+    m.add_class::<DType>()?;
+
     expr::init(py, m)?;
     signals::init(py, m)?;
     semantics::init(py, m)?;

pyargus/src/semantics.rs

@@ -7,8 +7,8 @@ use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyString};
 
 use crate::expr::PyBoolExpr;
-use crate::signals::{BoolSignal, FloatSignal, IntSignal, Kind, PySignal, UnsignedIntSignal};
+use crate::signals::{BoolSignal, FloatSignal, IntSignal, PySignal, UnsignedIntSignal};
-use crate::PyArgusError;
+use crate::{DType, PyArgusError};
 
 #[derive(Debug, Clone, derive_more::From, derive_more::TryInto)]
 #[try_into(owned, ref, ref_mut)]
@@ -42,16 +42,16 @@ impl PyTrace {
             })?;
             let kind = val.borrow().kind;
             let signal: SignalKind = match kind {
-                Kind::Bool => val.downcast::<PyCell<BoolSignal>>().unwrap().borrow().0.clone().into(),
+                DType::Bool => val.downcast::<PyCell<BoolSignal>>().unwrap().borrow().0.clone().into(),
-                Kind::Int => val.downcast::<PyCell<IntSignal>>().unwrap().borrow().0.clone().into(),
+                DType::Int => val.downcast::<PyCell<IntSignal>>().unwrap().borrow().0.clone().into(),
-                Kind::UnsignedInt => val
+                DType::UnsignedInt => val
                     .downcast::<PyCell<UnsignedIntSignal>>()
                     .unwrap()
                     .borrow()
                     .0
                     .clone()
                     .into(),
-                Kind::Float => val.downcast::<PyCell<FloatSignal>>().unwrap().borrow().0.clone().into(),
+                DType::Float => val.downcast::<PyCell<FloatSignal>>().unwrap().borrow().0.clone().into(),
             };
 
             signals.insert(key.to_string(), signal);

pyargus/src/signals.rs

@@ -3,20 +3,12 @@ use std::time::Duration;
 use argus_core::signals::{InterpolationMethod, Signal};
 use pyo3::prelude::*;
 
-use crate::PyArgusError;
+use crate::{DType, PyArgusError};
-
-#[derive(Copy, Clone, Debug)]
-pub enum Kind {
-    Bool,
-    Int,
-    UnsignedInt,
-    Float,
-}
 
 #[pyclass(name = "Signal", subclass, module = "argus")]
 #[derive(Debug, Clone)]
 pub struct PySignal {
-    pub kind: Kind,
+    pub kind: DType,
     pub interpolation: InterpolationMethod,
 }
 
@@ -32,7 +24,7 @@ macro_rules! impl_signals {
                 pub fn super_type() -> PySignal {
                     PySignal {
                         interpolation: InterpolationMethod::Linear,
-                        kind: Kind::$ty_name,
+                        kind: DType::$ty_name,
                     }
                 }
             }