add nlab cli to the python package

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]

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"]

src/python.rs

@@ -2,16 +2,29 @@ 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;
 
 #[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::<AnalogSignalPolarity>()?;
     m.add_class::<PowerState>()?;
     m.add_class::<PowerStatus>()?;
+    m.add_function(wrap_pyfunction!(run_cli, m)?)?;
     Ok(())
 }

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(())
+    }
 }

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,
+}