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