1use 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#[pyclass(name = "GetResult")]
24pub struct PyGetResult {
25 #[pyo3(get)]
27 pub pv_name: String,
28 value: PyObject,
29 #[pyo3(get)]
31 pub raw_pva: Vec<u8>,
32 #[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 #[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#[pyclass(name = "DiscoveredServer")]
71#[derive(Clone)]
72pub struct PyDiscoveredServer {
73 #[pyo3(get)]
75 pub guid: Vec<u8>,
76 #[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#[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 fn port(mut slf: PyRefMut<'_, Self>, port: u16) -> PyRefMut<'_, Self> {
128 slf.tcp_port = port;
129 slf
130 }
131
132 fn udp_port(mut slf: PyRefMut<'_, Self>, port: u16) -> PyRefMut<'_, Self> {
134 slf.udp_port = port;
135 slf
136 }
137
138 fn timeout(mut slf: PyRefMut<'_, Self>, secs: f64) -> PyRefMut<'_, Self> {
140 slf.timeout_secs = secs;
141 slf
142 }
143
144 fn no_broadcast(mut slf: PyRefMut<'_, Self>, enabled: bool) -> PyRefMut<'_, Self> {
147 slf.no_broadcast = enabled;
148 slf
149 }
150
151 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 fn authnz_user(mut slf: PyRefMut<'_, Self>, user: String) -> PyRefMut<'_, Self> {
163 slf.authnz_user = Some(user);
164 slf
165 }
166
167 fn authnz_host(mut slf: PyRefMut<'_, Self>, host: String) -> PyRefMut<'_, Self> {
169 slf.authnz_host = Some(host);
170 slf
171 }
172
173 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 fn search_addr(mut slf: PyRefMut<'_, Self>, addr: String) -> PyRefMut<'_, Self> {
186 slf.search_addr = Some(addr);
187 slf
188 }
189
190 fn bind_addr(mut slf: PyRefMut<'_, Self>, addr: String) -> PyRefMut<'_, Self> {
193 slf.bind_addr = Some(addr);
194 slf
195 }
196
197 fn debug(mut slf: PyRefMut<'_, Self>, enabled: bool) -> PyRefMut<'_, Self> {
199 slf.debug = enabled;
200 slf
201 }
202
203 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#[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 #[staticmethod]
278 fn builder() -> PyClientBuilder {
279 PyClientBuilder::new()
280 }
281
282 #[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 #[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 #[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 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 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 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 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 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 #[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 *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
490struct SubscriptionState {
493 active: AtomicBool,
494 error: Mutex<Option<String>>,
495}
496
497#[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 #[getter]
517 fn pv_name(&self) -> &str {
518 &self.pv_name
519 }
520
521 #[getter]
523 fn is_active(&self) -> bool {
524 self.state.active.load(Ordering::SeqCst)
525 }
526
527 #[getter]
529 fn error(&self) -> Option<String> {
530 self.state.error.lock().unwrap().clone()
531 }
532
533 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 fn drop(&mut self) {
565 self.abort_task();
566 }
567}
568
569#[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}