Skip to main content

spvirit/
nt.rs

1//! Python wrappers for Normative Type structs.
2
3use pyo3::prelude::*;
4use pyo3::types::{PyDict, PyList};
5use spvirit_types::{
6    NtAlarm, NtControl, NtDisplay, NtNdArray, NtPayload, NtScalar, NtScalarArray, NtTable,
7    NtTimeStamp, PvValue,
8};
9
10use crate::convert::{py_to_scalar, py_to_scalar_array, scalar_array_to_py, scalar_to_py};
11
12// ─── PyAlarm ─────────────────────────────────────────────────────────────────
13
14#[pyclass(name = "Alarm")]
15#[derive(Clone)]
16pub struct PyAlarm {
17    #[pyo3(get)]
18    pub severity: i32,
19    #[pyo3(get)]
20    pub status: i32,
21    #[pyo3(get)]
22    pub message: String,
23}
24
25impl From<&NtAlarm> for PyAlarm {
26    fn from(a: &NtAlarm) -> Self {
27        Self {
28            severity: a.severity,
29            status: a.status,
30            message: a.message.clone(),
31        }
32    }
33}
34
35#[pymethods]
36impl PyAlarm {
37    #[new]
38    #[pyo3(signature = (severity=0, status=0, message=String::new()))]
39    fn py_new(severity: i32, status: i32, message: String) -> Self {
40        Self { severity, status, message }
41    }
42
43    fn __repr__(&self) -> String {
44        format!(
45            "Alarm(severity={}, status={}, message={:?})",
46            self.severity, self.status, self.message
47        )
48    }
49}
50
51// ─── PyTimeStamp ─────────────────────────────────────────────────────────────
52
53#[pyclass(name = "TimeStamp")]
54#[derive(Clone)]
55pub struct PyTimeStamp {
56    #[pyo3(get)]
57    pub seconds_past_epoch: i64,
58    #[pyo3(get)]
59    pub nanoseconds: i32,
60    #[pyo3(get)]
61    pub user_tag: i32,
62}
63
64impl From<&NtTimeStamp> for PyTimeStamp {
65    fn from(ts: &NtTimeStamp) -> Self {
66        Self {
67            seconds_past_epoch: ts.seconds_past_epoch,
68            nanoseconds: ts.nanoseconds,
69            user_tag: ts.user_tag,
70        }
71    }
72}
73
74#[pymethods]
75impl PyTimeStamp {
76    #[new]
77    #[pyo3(signature = (seconds_past_epoch=0, nanoseconds=0, user_tag=0))]
78    fn py_new(seconds_past_epoch: i64, nanoseconds: i32, user_tag: i32) -> Self {
79        Self { seconds_past_epoch, nanoseconds, user_tag }
80    }
81
82    fn __repr__(&self) -> String {
83        format!(
84            "TimeStamp(seconds={}, ns={})",
85            self.seconds_past_epoch, self.nanoseconds
86        )
87    }
88}
89
90// ─── PyDisplay ───────────────────────────────────────────────────────────────
91
92#[pyclass(name = "Display")]
93#[derive(Clone)]
94pub struct PyDisplay {
95    #[pyo3(get)]
96    pub limit_low: f64,
97    #[pyo3(get)]
98    pub limit_high: f64,
99    #[pyo3(get)]
100    pub description: String,
101    #[pyo3(get)]
102    pub units: String,
103    #[pyo3(get)]
104    pub precision: i32,
105}
106
107impl From<&NtDisplay> for PyDisplay {
108    fn from(d: &NtDisplay) -> Self {
109        Self {
110            limit_low: d.limit_low,
111            limit_high: d.limit_high,
112            description: d.description.clone(),
113            units: d.units.clone(),
114            precision: d.precision,
115        }
116    }
117}
118
119#[pymethods]
120impl PyDisplay {
121    #[new]
122    #[pyo3(signature = (limit_low=0.0, limit_high=0.0, description=String::new(), units=String::new(), precision=0))]
123    fn py_new(limit_low: f64, limit_high: f64, description: String, units: String, precision: i32) -> Self {
124        Self { limit_low, limit_high, description, units, precision }
125    }
126
127    fn __repr__(&self) -> String {
128        format!(
129            "Display(low={}, high={}, units={:?})",
130            self.limit_low, self.limit_high, self.units
131        )
132    }
133}
134
135// ─── PyControl ───────────────────────────────────────────────────────────────
136
137#[pyclass(name = "Control")]
138#[derive(Clone)]
139pub struct PyControl {
140    #[pyo3(get)]
141    pub limit_low: f64,
142    #[pyo3(get)]
143    pub limit_high: f64,
144    #[pyo3(get)]
145    pub min_step: f64,
146}
147
148impl From<&NtControl> for PyControl {
149    fn from(c: &NtControl) -> Self {
150        Self {
151            limit_low: c.limit_low,
152            limit_high: c.limit_high,
153            min_step: c.min_step,
154        }
155    }
156}
157
158#[pymethods]
159impl PyControl {
160    #[new]
161    #[pyo3(signature = (limit_low=0.0, limit_high=0.0, min_step=0.0))]
162    fn py_new(limit_low: f64, limit_high: f64, min_step: f64) -> Self {
163        Self { limit_low, limit_high, min_step }
164    }
165
166    fn __repr__(&self) -> String {
167        format!(
168            "Control(low={}, high={}, min_step={})",
169            self.limit_low, self.limit_high, self.min_step
170        )
171    }
172}
173
174// ─── PyNtScalar ──────────────────────────────────────────────────────────────
175
176#[pyclass(name = "NtScalar")]
177pub struct PyNtScalar {
178    inner: NtScalar,
179}
180
181impl PyNtScalar {
182    pub fn new(inner: NtScalar) -> Self {
183        Self { inner }
184    }
185}
186
187#[pymethods]
188impl PyNtScalar {
189    /// Create an NtScalar from a Python value with optional metadata.
190    #[new]
191    #[pyo3(signature = (value, units=String::new(), display_low=0.0, display_high=0.0, display_description=String::new(), display_precision=0, control_low=0.0, control_high=0.0, control_min_step=0.0, alarm_severity=0, alarm_status=0, alarm_message=String::new()))]
192    fn py_new(
193        value: &Bound<'_, PyAny>,
194        units: String,
195        display_low: f64,
196        display_high: f64,
197        display_description: String,
198        display_precision: i32,
199        control_low: f64,
200        control_high: f64,
201        control_min_step: f64,
202        alarm_severity: i32,
203        alarm_status: i32,
204        alarm_message: String,
205    ) -> PyResult<Self> {
206        let sv = py_to_scalar(value)?;
207        let mut nt = NtScalar::from_value(sv);
208        nt.units = units;
209        nt.display_low = display_low;
210        nt.display_high = display_high;
211        nt.display_description = display_description;
212        nt.display_precision = display_precision;
213        nt.control_low = control_low;
214        nt.control_high = control_high;
215        nt.control_min_step = control_min_step;
216        nt.alarm_severity = alarm_severity;
217        nt.alarm_status = alarm_status;
218        nt.alarm_message = alarm_message;
219        Ok(Self { inner: nt })
220    }
221
222    #[getter]
223    fn value(&self, py: Python<'_>) -> PyObject {
224        scalar_to_py(py, &self.inner.value)
225    }
226
227    #[getter]
228    fn alarm_severity(&self) -> i32 {
229        self.inner.alarm_severity
230    }
231
232    #[getter]
233    fn alarm_status(&self) -> i32 {
234        self.inner.alarm_status
235    }
236
237    #[getter]
238    fn alarm_message(&self) -> &str {
239        &self.inner.alarm_message
240    }
241
242    #[getter]
243    fn units(&self) -> &str {
244        &self.inner.units
245    }
246
247    #[getter]
248    fn display_low(&self) -> f64 {
249        self.inner.display_low
250    }
251
252    #[getter]
253    fn display_high(&self) -> f64 {
254        self.inner.display_high
255    }
256
257    #[getter]
258    fn display_description(&self) -> &str {
259        &self.inner.display_description
260    }
261
262    #[getter]
263    fn display_precision(&self) -> i32 {
264        self.inner.display_precision
265    }
266
267    #[getter]
268    fn control_low(&self) -> f64 {
269        self.inner.control_low
270    }
271
272    #[getter]
273    fn control_high(&self) -> f64 {
274        self.inner.control_high
275    }
276
277    #[getter]
278    fn control_min_step(&self) -> f64 {
279        self.inner.control_min_step
280    }
281
282    fn __repr__(&self, py: Python<'_>) -> String {
283        let val = scalar_to_py(py, &self.inner.value);
284        format!("NtScalar(value={}, units={:?})", val, self.inner.units)
285    }
286}
287
288// ─── PyNtScalarArray ─────────────────────────────────────────────────────────
289
290#[pyclass(name = "NtScalarArray")]
291pub struct PyNtScalarArray {
292    inner: NtScalarArray,
293}
294
295impl PyNtScalarArray {
296    pub fn new(inner: NtScalarArray) -> Self {
297        Self { inner }
298    }
299}
300
301#[pymethods]
302impl PyNtScalarArray {
303    /// Create an NtScalarArray from a Python list/bytes.
304    #[new]
305    fn py_new(value: &Bound<'_, PyAny>) -> PyResult<Self> {
306        let arr = py_to_scalar_array(value)?;
307        Ok(Self { inner: NtScalarArray::from_value(arr) })
308    }
309
310    #[getter]
311    fn value(&self, py: Python<'_>) -> PyObject {
312        scalar_array_to_py(py, &self.inner.value)
313    }
314
315    #[getter]
316    fn alarm(&self) -> PyAlarm {
317        PyAlarm::from(&self.inner.alarm)
318    }
319
320    #[getter]
321    fn time_stamp(&self) -> PyTimeStamp {
322        PyTimeStamp::from(&self.inner.time_stamp)
323    }
324
325    #[getter]
326    fn display(&self) -> PyDisplay {
327        PyDisplay::from(&self.inner.display)
328    }
329
330    #[getter]
331    fn control(&self) -> PyControl {
332        PyControl::from(&self.inner.control)
333    }
334
335    fn __repr__(&self, py: Python<'_>) -> String {
336        let val = scalar_array_to_py(py, &self.inner.value);
337        format!("NtScalarArray(value={})", val)
338    }
339}
340
341// ─── PyNtTable ───────────────────────────────────────────────────────────────
342
343#[pyclass(name = "NtTable")]
344pub struct PyNtTable {
345    inner: NtTable,
346}
347
348impl PyNtTable {
349    pub fn new(inner: NtTable) -> Self {
350        Self { inner }
351    }
352}
353
354#[pymethods]
355impl PyNtTable {
356    #[getter]
357    fn labels(&self, py: Python<'_>) -> PyResult<PyObject> {
358        let items: Vec<PyObject> = self
359            .inner
360            .labels
361            .iter()
362            .map(|s| pyo3::types::PyString::new(py, s).into_any().unbind())
363            .collect();
364        Ok(PyList::new(py, &items)?.into_any().unbind())
365    }
366
367    /// Return columns as a dict of {name: list}.
368    fn columns(&self, py: Python<'_>) -> PyResult<PyObject> {
369        let dict = pyo3::types::PyDict::new(py);
370        for col in &self.inner.columns {
371            dict.set_item(&col.name, scalar_array_to_py(py, &col.values))?;
372        }
373        Ok(dict.into_any().unbind())
374    }
375
376    #[getter]
377    fn descriptor(&self) -> Option<&str> {
378        self.inner.descriptor.as_deref()
379    }
380
381    #[getter]
382    fn alarm(&self) -> Option<PyAlarm> {
383        self.inner.alarm.as_ref().map(PyAlarm::from)
384    }
385
386    #[getter]
387    fn time_stamp(&self) -> Option<PyTimeStamp> {
388        self.inner.time_stamp.as_ref().map(PyTimeStamp::from)
389    }
390
391    fn __repr__(&self) -> String {
392        format!(
393            "NtTable(labels={:?}, columns={})",
394            self.inner.labels,
395            self.inner.columns.len()
396        )
397    }
398}
399
400// ─── PyNtNdArray ─────────────────────────────────────────────────────────────
401
402#[pyclass(name = "NtNdArray")]
403pub struct PyNtNdArray {
404    inner: NtNdArray,
405}
406
407impl PyNtNdArray {
408    pub fn new(inner: NtNdArray) -> Self {
409        Self { inner }
410    }
411}
412
413#[pymethods]
414impl PyNtNdArray {
415    #[getter]
416    fn value(&self, py: Python<'_>) -> PyObject {
417        scalar_array_to_py(py, &self.inner.value)
418    }
419
420    #[getter]
421    fn unique_id(&self) -> i32 {
422        self.inner.unique_id
423    }
424
425    #[getter]
426    fn compressed_size(&self) -> i64 {
427        self.inner.compressed_size
428    }
429
430    #[getter]
431    fn uncompressed_size(&self) -> i64 {
432        self.inner.uncompressed_size
433    }
434
435    /// Return dimensions as a list of dicts.
436    fn dimensions(&self, py: Python<'_>) -> PyResult<PyObject> {
437        let items: Vec<PyObject> = self
438            .inner
439            .dimension
440            .iter()
441            .map(|d| {
442                let dict = pyo3::types::PyDict::new(py);
443                dict.set_item("size", d.size).expect("set");
444                dict.set_item("offset", d.offset).expect("set");
445                dict.set_item("full_size", d.full_size).expect("set");
446                dict.set_item("binning", d.binning).expect("set");
447                dict.set_item("reverse", d.reverse).expect("set");
448                dict.into_any().unbind()
449            })
450            .collect();
451        Ok(PyList::new(py, &items)?.into_any().unbind())
452    }
453
454    #[getter]
455    fn data_time_stamp(&self) -> PyTimeStamp {
456        PyTimeStamp::from(&self.inner.data_time_stamp)
457    }
458
459    fn __repr__(&self) -> String {
460        let dims: Vec<i32> = self.inner.dimension.iter().map(|d| d.size).collect();
461        format!("NtNdArray(unique_id={}, dims={:?})", self.inner.unique_id, dims)
462    }
463}
464
465// ─── NtPayload → Python wrapper ──────────────────────────────────────────────
466
467/// Extract the Rust `NtPayload` from a Python NT object.
468pub fn py_to_nt_payload(obj: &Bound<'_, PyAny>) -> PyResult<NtPayload> {
469    if let Ok(s) = obj.downcast::<PyNtScalar>() {
470        Ok(NtPayload::Scalar(s.borrow().inner.clone()))
471    } else if let Ok(a) = obj.downcast::<PyNtScalarArray>() {
472        Ok(NtPayload::ScalarArray(a.borrow().inner.clone()))
473    } else if let Ok(t) = obj.downcast::<PyNtTable>() {
474        Ok(NtPayload::Table(t.borrow().inner.clone()))
475    } else if let Ok(n) = obj.downcast::<PyNtNdArray>() {
476        Ok(NtPayload::NdArray(n.borrow().inner.clone()))
477    } else {
478        Err(pyo3::exceptions::PyTypeError::new_err(
479            "expected NtScalar, NtScalarArray, NtTable, or NtNdArray",
480        ))
481    }
482}
483
484pub fn nt_payload_to_py(py: Python<'_>, payload: NtPayload) -> PyObject {
485    match payload {
486        NtPayload::Scalar(s) => PyNtScalar::new(s).into_pyobject(py).expect("NtScalar").into_any().unbind(),
487        NtPayload::ScalarArray(a) => PyNtScalarArray::new(a).into_pyobject(py).expect("NtScalarArray").into_any().unbind(),
488        NtPayload::Table(t) => PyNtTable::new(t).into_pyobject(py).expect("NtTable").into_any().unbind(),
489        NtPayload::NdArray(n) => PyNtNdArray::new(n).into_pyobject(py).expect("NtNdArray").into_any().unbind(),
490        NtPayload::Enum(e) => {
491            let d = PyDict::new(py);
492            d.set_item("index", e.index).ok();
493            d.set_item("choices", &e.choices).ok();
494            d.set_item("selected", e.selected().unwrap_or("")).ok();
495            d.unbind().into_any()
496        }
497        NtPayload::Generic { struct_id, fields } => {
498            let d = PyDict::new(py);
499            d.set_item("struct_id", &struct_id).ok();
500            for (name, val) in &fields {
501                d.set_item(name, pvvalue_to_py(py, val)).ok();
502            }
503            d.unbind().into_any()
504        }
505    }
506}
507
508fn pvvalue_to_py(py: Python<'_>, val: &PvValue) -> PyObject {
509    match val {
510        PvValue::Scalar(s) => scalar_to_py(py, s),
511        PvValue::ScalarArray(a) => scalar_array_to_py(py, a),
512        PvValue::Structure { struct_id, fields } => {
513            let d = PyDict::new(py);
514            d.set_item("struct_id", struct_id).ok();
515            for (name, v) in fields {
516                d.set_item(name, pvvalue_to_py(py, v)).ok();
517            }
518            d.unbind().into_any()
519        }
520    }
521}