Skip to main content

spvirit/
client.rs

1//! Python client wrappers — sync-only for phase 1.
2
3use std::net::SocketAddr;
4use std::ops::ControlFlow;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{Arc, Mutex};
7use std::time::Duration;
8
9use pyo3::prelude::*;
10use pyo3::types::{PyDict, PyList};
11
12use spvirit_client::pva_client::PvaClient;
13use spvirit_client::search::{build_auto_broadcast_targets, discover_servers};
14
15use crate::convert::{decoded_to_py, py_to_json};
16use crate::errors::to_py_err;
17use crate::runtime::RUNTIME;
18
19// ─── GetResult ───────────────────────────────────────────────────────────────
20
21/// Result of a get operation: `.pv_name`, `.value` (decoded Python value),
22/// `.raw_pva` (raw PVA frame bytes), `.raw_pvd` (raw pvData body bytes).
23#[pyclass(name = "GetResult")]
24pub struct PyGetResult {
25    /// Name of the PV this result was read from.
26    #[pyo3(get)]
27    pub pv_name: String,
28    value: PyObject,
29    /// Raw PVA frame bytes of the get response.
30    #[pyo3(get)]
31    pub raw_pva: Vec<u8>,
32    /// Raw pvData body bytes of the get response.
33    #[pyo3(get)]
34    pub raw_pvd: Vec<u8>,
35}
36
37impl PyGetResult {
38    pub(crate) fn new(
39        pv_name: String,
40        value: PyObject,
41        raw_pva: Vec<u8>,
42        raw_pvd: Vec<u8>,
43    ) -> Self {
44        Self {
45            pv_name,
46            value,
47            raw_pva,
48            raw_pvd,
49        }
50    }
51}
52
53#[pymethods]
54impl PyGetResult {
55    /// Decoded Python value (usually a dict mirroring the NT structure).
56    #[getter]
57    fn value(&self, py: Python<'_>) -> PyObject {
58        self.value.clone_ref(py)
59    }
60
61    fn __repr__(&self) -> String {
62        format!("GetResult(pv_name={:?})", self.pv_name)
63    }
64}
65
66// ─── DiscoveredServer ────────────────────────────────────────────────────────
67
68/// A server found by `discover_servers()`: `.guid` (12-byte bytes) and
69/// `.tcp_addr` (`"ip:port"`).
70#[pyclass(name = "DiscoveredServer")]
71#[derive(Clone)]
72pub struct PyDiscoveredServer {
73    /// Server GUID (12 bytes).
74    #[pyo3(get)]
75    pub guid: Vec<u8>,
76    /// Server TCP address as `"ip:port"`.
77    #[pyo3(get)]
78    pub tcp_addr: String,
79}
80
81#[pymethods]
82impl PyDiscoveredServer {
83    fn __repr__(&self) -> String {
84        format!("DiscoveredServer(tcp_addr={:?})", self.tcp_addr)
85    }
86}
87
88// ─── ClientBuilder ───────────────────────────────────────────────────────────
89
90/// Builder for `Client`. Defaults: TCP 5075, UDP 5076, timeout 5.0 s,
91/// broadcast search enabled.
92#[pyclass(name = "ClientBuilder")]
93pub struct PyClientBuilder {
94    udp_port: u16,
95    tcp_port: u16,
96    timeout_secs: f64,
97    no_broadcast: bool,
98    name_servers: Vec<String>,
99    authnz_user: Option<String>,
100    authnz_host: Option<String>,
101    server_addr: Option<String>,
102    search_addr: Option<String>,
103    bind_addr: Option<String>,
104    debug: bool,
105}
106
107#[pymethods]
108impl PyClientBuilder {
109    #[new]
110    fn new() -> Self {
111        Self {
112            udp_port: 5076,
113            tcp_port: 5075,
114            timeout_secs: 5.0,
115            no_broadcast: false,
116            name_servers: Vec::new(),
117            authnz_user: None,
118            authnz_host: None,
119            server_addr: None,
120            search_addr: None,
121            bind_addr: None,
122            debug: false,
123        }
124    }
125
126    /// Set the default TCP port used when connecting to servers.
127    fn port(mut slf: PyRefMut<'_, Self>, port: u16) -> PyRefMut<'_, Self> {
128        slf.tcp_port = port;
129        slf
130    }
131
132    /// Set the UDP port used for broadcast search.
133    fn udp_port(mut slf: PyRefMut<'_, Self>, port: u16) -> PyRefMut<'_, Self> {
134        slf.udp_port = port;
135        slf
136    }
137
138    /// Set the operation timeout in seconds.
139    fn timeout(mut slf: PyRefMut<'_, Self>, secs: f64) -> PyRefMut<'_, Self> {
140        slf.timeout_secs = secs;
141        slf
142    }
143
144    /// Disable UDP broadcast search when `enabled` is True (use name servers
145    /// or a fixed server address instead).
146    fn no_broadcast(mut slf: PyRefMut<'_, Self>, enabled: bool) -> PyRefMut<'_, Self> {
147        slf.no_broadcast = enabled;
148        slf
149    }
150
151    /// Add a name server `"ip:port"` address to query for PV names.
152    /// Raises ValueError on an invalid address.
153    fn name_server(mut slf: PyRefMut<'_, Self>, addr: String) -> PyResult<PyRefMut<'_, Self>> {
154        let _: SocketAddr = addr.parse().map_err(|e| {
155            pyo3::exceptions::PyValueError::new_err(format!("invalid address: {e}"))
156        })?;
157        slf.name_servers.push(addr);
158        Ok(slf)
159    }
160
161    /// Set the user name reported during connection authentication.
162    fn authnz_user(mut slf: PyRefMut<'_, Self>, user: String) -> PyRefMut<'_, Self> {
163        slf.authnz_user = Some(user);
164        slf
165    }
166
167    /// Set the host name reported during connection authentication.
168    fn authnz_host(mut slf: PyRefMut<'_, Self>, host: String) -> PyRefMut<'_, Self> {
169        slf.authnz_host = Some(host);
170        slf
171    }
172
173    /// Connect directly to the server at `"ip:port"`, skipping search.
174    /// Raises ValueError on an invalid address.
175    fn server_addr(mut slf: PyRefMut<'_, Self>, addr: String) -> PyResult<PyRefMut<'_, Self>> {
176        let _: SocketAddr = addr.parse().map_err(|e| {
177            pyo3::exceptions::PyValueError::new_err(format!("invalid address: {e}"))
178        })?;
179        slf.server_addr = Some(addr);
180        Ok(slf)
181    }
182
183    /// Set the IP address search requests are sent to (validated at
184    /// `build()`; ValueError on a bad IP).
185    fn search_addr(mut slf: PyRefMut<'_, Self>, addr: String) -> PyRefMut<'_, Self> {
186        slf.search_addr = Some(addr);
187        slf
188    }
189
190    /// Set the local IP address to bind the search socket to (validated at
191    /// `build()`; ValueError on a bad IP).
192    fn bind_addr(mut slf: PyRefMut<'_, Self>, addr: String) -> PyRefMut<'_, Self> {
193        slf.bind_addr = Some(addr);
194        slf
195    }
196
197    /// Enable verbose protocol debug logging.
198    fn debug(mut slf: PyRefMut<'_, Self>, enabled: bool) -> PyRefMut<'_, Self> {
199        slf.debug = enabled;
200        slf
201    }
202
203    /// Build and return a configured `Client`. Raises ValueError if any
204    /// stored address string fails to parse.
205    fn build(&self) -> PyResult<PyClient> {
206        let mut b = PvaClient::builder()
207            .port(self.tcp_port)
208            .udp_port(self.udp_port)
209            .timeout(Duration::from_secs_f64(self.timeout_secs));
210        if self.no_broadcast {
211            b = b.no_broadcast();
212        }
213        for ns in &self.name_servers {
214            let addr: SocketAddr = ns.parse().map_err(|e| {
215                pyo3::exceptions::PyValueError::new_err(format!("invalid address: {e}"))
216            })?;
217            b = b.name_server(addr);
218        }
219        if let Some(ref user) = self.authnz_user {
220            b = b.authnz_user(user);
221        }
222        if let Some(ref host) = self.authnz_host {
223            b = b.authnz_host(host);
224        }
225        if let Some(ref addr) = self.server_addr {
226            let sa: SocketAddr = addr.parse().map_err(|e| {
227                pyo3::exceptions::PyValueError::new_err(format!("invalid address: {e}"))
228            })?;
229            b = b.server_addr(sa);
230        }
231        if let Some(ref addr) = self.search_addr {
232            let ip: std::net::IpAddr = addr
233                .parse()
234                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("invalid IP: {e}")))?;
235            b = b.search_addr(ip);
236        }
237        if let Some(ref addr) = self.bind_addr {
238            let ip: std::net::IpAddr = addr
239                .parse()
240                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("invalid IP: {e}")))?;
241            b = b.bind_addr(ip);
242        }
243        if self.debug {
244            b = b.debug();
245        }
246        Ok(PyClient { inner: b.build() })
247    }
248
249    fn __repr__(&self) -> String {
250        format!(
251            "<spvirit.ClientBuilder (tcp={}, udp={}, timeout={}s)>",
252            self.tcp_port, self.udp_port, self.timeout_secs
253        )
254    }
255}
256
257// ─── Client ──────────────────────────────────────────────────────────────────
258
259/// High-level PVAccess client: get/put/monitor/subscribe/info/pvlist.
260/// `Client()` uses broadcast-search defaults; `Client.builder()` configures.
261/// Operations raise the SpviritError hierarchy.
262#[pyclass(name = "Client")]
263pub struct PyClient {
264    inner: PvaClient,
265}
266
267#[pymethods]
268impl PyClient {
269    #[new]
270    fn new() -> Self {
271        Self {
272            inner: PvaClient::builder().build(),
273        }
274    }
275
276    /// Create a builder for fine-grained configuration.
277    #[staticmethod]
278    fn builder() -> PyClientBuilder {
279        PyClientBuilder::new()
280    }
281
282    /// Fetch the current value of a PV (blocking).
283    ///
284    /// If `fields` is provided (a list of dotted paths or a single string,
285    /// e.g. `["value", "alarm.severity"]`), the pvRequest restricts the
286    /// returned structure to those paths.
287    #[pyo3(signature = (pv_name, fields=None))]
288    fn get(
289        &self,
290        py: Python<'_>,
291        pv_name: String,
292        fields: Option<PyObject>,
293    ) -> PyResult<PyGetResult> {
294        let client = self.inner.clone();
295        let fields = crate::channel::normalize_fields(py, fields)?;
296        let result = crate::runtime::block_on_py(py, async {
297            if fields.is_empty() {
298                client.pvget(&pv_name).await
299            } else {
300                let refs: Vec<&str> = fields.iter().map(String::as_str).collect();
301                client.pvget_fields(&pv_name, &refs).await
302            }
303        })
304        .map_err(to_py_err)?;
305        let value = decoded_to_py(py, &result.value);
306        Ok(PyGetResult {
307            pv_name: result.pv_name,
308            value,
309            raw_pva: result.raw_pva,
310            raw_pvd: result.raw_pvd,
311        })
312    }
313
314    /// Write a value to a PV (blocking).
315    ///
316    /// `fields` selects which pvRequest fields are targeted (a list of
317    /// dotted paths or a single string). Defaults to `["value"]` when
318    /// omitted.
319    #[pyo3(signature = (pv_name, value, fields=None))]
320    fn put(
321        &self,
322        py: Python<'_>,
323        pv_name: String,
324        value: PyObject,
325        fields: Option<PyObject>,
326    ) -> PyResult<()> {
327        let json_val = py_to_json(value.bind(py))?;
328        let client = self.inner.clone();
329        let fields = crate::channel::normalize_fields(py, fields)?;
330        crate::runtime::block_on_py(py, async {
331            if fields.is_empty() {
332                client.pvput(&pv_name, json_val).await
333            } else {
334                let refs: Vec<&str> = fields.iter().map(String::as_str).collect();
335                client.pvput_fields(&pv_name, json_val, &refs).await
336            }
337        })
338        .map_err(to_py_err)
339    }
340
341    /// Subscribe to a PV and call `callback(value_dict)` for each update.
342    ///
343    /// Blocks (GIL released between updates) until the callback returns
344    /// `False` or raises — a raised exception stops the monitor and
345    /// propagates to the caller. `fields` restricts the subscription to the
346    /// given dotted paths. For a non-blocking variant see `subscribe`.
347    #[pyo3(signature = (pv_name, callback, fields=None))]
348    fn monitor(
349        &self,
350        py: Python<'_>,
351        pv_name: String,
352        callback: PyObject,
353        fields: Option<PyObject>,
354    ) -> PyResult<()> {
355        let client = self.inner.clone();
356        let fields = crate::channel::normalize_fields(py, fields)?;
357        let mut cb_err: Option<PyErr> = None;
358        // The GIL is released for the wait; the callback reacquires it per
359        // update via with_gil.
360        let result = crate::runtime::block_on_py(py, {
361            let cb_err = &mut cb_err;
362            async move {
363                let refs: Vec<&str> = fields.iter().map(String::as_str).collect();
364                client
365                    .pvmonitor_fields(&pv_name, &refs, |decoded| {
366                        let keep_going = Python::with_gil(|py| {
367                            let py_val = decoded_to_py(py, decoded);
368                            match callback.call1(py, (py_val,)) {
369                                Ok(ret) => {
370                                    // If callback returns False, stop
371                                    ret.extract::<bool>(py).unwrap_or(true)
372                                }
373                                Err(e) => {
374                                    *cb_err = Some(e);
375                                    false
376                                }
377                            }
378                        });
379                        if keep_going {
380                            ControlFlow::Continue(())
381                        } else {
382                            ControlFlow::Break(())
383                        }
384                    })
385                    .await
386            }
387        });
388        if let Some(e) = cb_err {
389            return Err(e);
390        }
391        result.map_err(to_py_err)
392    }
393
394    /// Retrieve introspection (field description) for a PV.
395    fn info(&self, py: Python<'_>, pv_name: String) -> PyResult<PyObject> {
396        let client = self.inner.clone();
397        let desc = crate::runtime::block_on_py(py, client.pvinfo(&pv_name)).map_err(to_py_err)?;
398        // Return as a dict: {struct_id, fields: [{name, field_type}, ...]}
399        let dict = PyDict::new(py);
400        dict.set_item("struct_id", &desc.struct_id)?;
401        let fields: Vec<PyObject> = desc
402            .fields
403            .iter()
404            .map(|f| {
405                let fd = PyDict::new(py);
406                fd.set_item("name", &f.name).expect("set");
407                fd.set_item("field_type", format!("{:?}", f.field_type))
408                    .expect("set");
409                fd.into_any().unbind()
410            })
411            .collect();
412        dict.set_item("fields", PyList::new(py, &fields)?)?;
413        Ok(dict.into_any().unbind())
414    }
415
416    /// List PV names from a specific server.
417    fn pvlist(&self, py: Python<'_>, server_addr: String) -> PyResult<Vec<String>> {
418        let addr: SocketAddr = server_addr.parse().map_err(|e| {
419            pyo3::exceptions::PyValueError::new_err(format!("invalid address: {e}"))
420        })?;
421        let client = self.inner.clone();
422        crate::runtime::block_on_py(py, client.pvlist(addr)).map_err(to_py_err)
423    }
424
425    /// Subscribe to a PV without blocking; returns a `Subscription` handle.
426    ///
427    /// `callback(value)` runs on a background runtime thread for each update,
428    /// sequentially per subscription. Returning `False` from the callback
429    /// unsubscribes, matching `monitor`; raising also unsubscribes and stores
430    /// the message in `subscription.error`. Call `subscription.close()` to
431    /// stop promptly — it works even while the PV is quiet. If the
432    /// subscription ends on a network/protocol error, `subscription.error`
433    /// holds the message and `is_active` becomes `False`.
434    #[pyo3(signature = (pv_name, callback, fields=None))]
435    fn subscribe(
436        &self,
437        py: Python<'_>,
438        pv_name: String,
439        callback: PyObject,
440        fields: Option<PyObject>,
441    ) -> PyResult<PySubscription> {
442        let client = self.inner.clone();
443        let fields = crate::channel::normalize_fields(py, fields)?;
444        let state = Arc::new(SubscriptionState {
445            active: AtomicBool::new(true),
446            error: Mutex::new(None),
447        });
448        let task_state = state.clone();
449        let task_pv = pv_name.clone();
450        let handle = RUNTIME.spawn(async move {
451            let refs: Vec<&str> = fields.iter().map(String::as_str).collect();
452            let result = client
453                .pvmonitor_fields(&task_pv, &refs, |decoded| {
454                    let keep_going = Python::with_gil(|py| {
455                        let py_val = decoded_to_py(py, decoded);
456                        match callback.call1(py, (py_val,)) {
457                            Ok(ret) => ret.extract::<bool>(py).unwrap_or(true),
458                            Err(e) => {
459                                // No caller to raise into: record the failure
460                                // on the subscription and stop.
461                                *task_state.error.lock().unwrap() = Some(e.to_string());
462                                false
463                            }
464                        }
465                    });
466                    if keep_going {
467                        ControlFlow::Continue(())
468                    } else {
469                        ControlFlow::Break(())
470                    }
471                })
472                .await;
473            if let Err(e) = result {
474                *task_state.error.lock().unwrap() = Some(e.to_string());
475            }
476            task_state.active.store(false, Ordering::SeqCst);
477        });
478        Ok(PySubscription {
479            pv_name,
480            state,
481            handle: Mutex::new(Some(handle)),
482        })
483    }
484
485    fn __repr__(&self) -> &'static str {
486        "<spvirit.Client>"
487    }
488}
489
490// ─── Subscription ────────────────────────────────────────────────────────────
491
492struct SubscriptionState {
493    active: AtomicBool,
494    error: Mutex<Option<String>>,
495}
496
497/// Handle to a non-blocking monitor started with `Client.subscribe`.
498#[pyclass(name = "Subscription")]
499pub struct PySubscription {
500    pv_name: String,
501    state: Arc<SubscriptionState>,
502    handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
503}
504
505impl PySubscription {
506    fn abort_task(&self) {
507        if let Some(h) = self.handle.lock().unwrap().take() {
508            h.abort();
509        }
510    }
511}
512
513#[pymethods]
514impl PySubscription {
515    /// Name of the PV this subscription watches.
516    #[getter]
517    fn pv_name(&self) -> &str {
518        &self.pv_name
519    }
520
521    /// True while updates are still being delivered.
522    #[getter]
523    fn is_active(&self) -> bool {
524        self.state.active.load(Ordering::SeqCst)
525    }
526
527    /// Error message if the subscription ended on a failure, else None.
528    #[getter]
529    fn error(&self) -> Option<String> {
530        self.state.error.lock().unwrap().clone()
531    }
532
533    /// Stop the subscription. Idempotent; returns immediately and works
534    /// even while the PV is quiet.
535    fn close(&self) {
536        self.abort_task();
537        self.state.active.store(false, Ordering::SeqCst);
538    }
539
540    fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
541        slf
542    }
543
544    #[pyo3(signature = (_exc_type=None, _exc=None, _tb=None))]
545    fn __exit__(
546        &self,
547        _exc_type: Option<PyObject>,
548        _exc: Option<PyObject>,
549        _tb: Option<PyObject>,
550    ) -> bool {
551        self.close();
552        false
553    }
554
555    fn __repr__(&self) -> String {
556        let status = if self.is_active() { "active" } else { "closed" };
557        format!("<spvirit.Subscription '{}' ({status})>", self.pv_name)
558    }
559}
560
561impl Drop for PySubscription {
562    // A subscription nobody holds a handle to is uncontrollable; stop it
563    // rather than leak a detached task.
564    fn drop(&mut self) {
565        self.abort_task();
566    }
567}
568
569// ─── discover_servers ────────────────────────────────────────────────────────
570
571/// Discover PVA servers on the network via UDP broadcast search.
572///
573/// `udp_port` (default 5076) is the search port, `timeout` (default 2.0 s)
574/// how long to collect responses, `debug` enables verbose logging.
575/// Returns a list of `DiscoveredServer`.
576#[pyfunction]
577#[pyo3(signature = (udp_port=5076, timeout=2.0, debug=false))]
578pub fn py_discover_servers(
579    py: Python<'_>,
580    udp_port: u16,
581    timeout: f64,
582    debug: bool,
583) -> PyResult<Vec<PyDiscoveredServer>> {
584    let targets = build_auto_broadcast_targets();
585    let dur = Duration::from_secs_f64(timeout);
586    let servers = crate::runtime::block_on_py(py, discover_servers(udp_port, dur, &targets, debug))
587        .map_err(to_py_err)?;
588    Ok(servers
589        .into_iter()
590        .map(|s| PyDiscoveredServer {
591            guid: s.guid.to_vec(),
592            tcp_addr: s.tcp_addr.to_string(),
593        })
594        .collect())
595}