diff --git a/Cargo.toml b/Cargo.toml index 7bb64e2..8aead9d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ log = "~0.4" regex = "~1" rusb = { version = "0.9.3", features = ["vendored"] } dfu-libusb = "0.5.1" +clap = { version = "4.5.16", features = ["derive"] } pyo3 = { version = "~0.22", features = ["multiple-pymethods"] } [dev-dependencies] diff --git a/pyproject.toml b/pyproject.toml index 28b1de9..1ebdacf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,12 +5,13 @@ build-backend = "maturin" [project] name = "nlabapi" requires-python = ">=3.8" -classifiers = [ - "Programming Language :: Rust", - "Programming Language :: Python :: Implementation :: CPython", - "Programming Language :: Python :: Implementation :: PyPy", -] +readme = "README.md" +license = { file = "LICENSE" } dynamic = ["version"] + +[project.scripts] +nlab = "nlabapi:run_cli" + [tool.maturin] module-name = "nlabapi" features = ["pyo3/extension-module"] diff --git a/src/python.rs b/src/python.rs index 5f73f10..2cc34ba 100644 --- a/src/python.rs +++ b/src/python.rs @@ -2,9 +2,12 @@ mod bench; mod scope; mod analog_output; mod pulse_output; +mod cli; use pyo3::prelude::*; use crate::{AnalogSignalPolarity, AnalogWaveType, PowerStatus, PowerState}; +use cli::{Cli, Commands}; +use clap::Parser; #[pyclass] struct LabBench; @@ -12,6 +15,16 @@ struct LabBench; #[pyclass] struct Nlab(crate::Nlab); +#[pyfunction] +fn run_cli(_py: Python) -> PyResult<()> { + let args: Vec<_> = std::env::args_os().skip(1).collect(); + let cli = Cli::parse_from(args); + + match &cli.command { + Commands::Update => { LabBench::update_all_nlabs() } + } +} + /// A Python module implemented in Rust. #[pymodule] fn nlabapi(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -21,5 +34,6 @@ fn nlabapi(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_function(wrap_pyfunction!(run_cli, m)?)?; Ok(()) } \ No newline at end of file diff --git a/src/python/bench.rs b/src/python/bench.rs index 4fb571f..72bccac 100644 --- a/src/python/bench.rs +++ b/src/python/bench.rs @@ -1,3 +1,5 @@ +use std::thread; +use std::time::Duration; use pyo3::exceptions::*; use pyo3::prelude::*; use crate::{LabBench, python}; @@ -25,4 +27,40 @@ impl python::LabBench { println!("Cannot create LabBench"); } } + + #[staticmethod] + fn count_connected_nlabs() -> usize { + if let Ok(bench) = LabBench::new() { + return bench.list().count(); + } + 0 + } + + #[staticmethod] + pub(super) fn update_all_nlabs() -> PyResult<()> { + if let Ok(mut bench) = LabBench::new() { + for nlab_link in bench.list() { + // Request DFU on any nLab that is available + if nlab_link.available { + if let Err(e) = nlab_link.request_dfu() { + println!("Failed to request DFU on an available nLab: {e}"); + return Err(PyRuntimeError::new_err(format!("{e}"))); + } + } + } + println!("Updating all connected nLabs..."); + // Wait 500ms for the scope to detach and re-attach as DFU + thread::sleep(Duration::from_millis(500)); + bench.refresh(); + + for nlab_link in bench.list() { + if let Err(e) = nlab_link.update() { + println!("Encountered an error updating nLab: {e}"); + return Err(PyRuntimeError::new_err(format!("{e}"))); + } + } + println!("Update complete!"); + } + Ok(()) + } } diff --git a/src/python/cli.rs b/src/python/cli.rs new file mode 100644 index 0000000..fcdbada --- /dev/null +++ b/src/python/cli.rs @@ -0,0 +1,15 @@ +use clap::{Parser, Subcommand}; + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +#[command(propagate_version = true)] +pub(super) struct Cli { + #[command(subcommand)] + pub command: Commands, +} + +#[derive(Subcommand, Debug)] +pub(super) enum Commands { + /// Update all detected nLabs + Update, +} \ No newline at end of file