Ship Vortex CLI as part of Python package (#6233)

After installing the python package, the CLI can be run as `vx`.

Fixes #2589

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

Cargo.lock

@@ -285,6 +285,7 @@ vortex-cuda-macros = { path = "./vortex-cuda/macros" }
 vortex-duckdb = { path = "./vortex-duckdb", default-features = false }
 vortex-test-e2e = { path = "./vortex-test/e2e", default-features = false }
 vortex-test-e2e-cuda = { path = "./vortex-test/e2e-cuda", default-features = false }
+vortex-tui = { path = "./vortex-tui" }
 
 [workspace.dependencies.getrandom_v03]
 features = ["wasm_js"]

Cargo.toml

@@ -52,6 +52,7 @@ tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
 tracing = { workspace = true, features = ["std", "log"] }
 url = { workspace = true }
 vortex = { workspace = true, features = ["object_store", "tokio"] }
+vortex-tui = { workspace = true }
 
 [dev-dependencies]
 vortex-array = { workspace = true, features = ["_test-harness"] }

vortex-python/Cargo.toml

@@ -33,6 +33,9 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 
+[project.scripts]
+vx = "vortex.cli:main"
+
 [project.optional-dependencies]
 polars = ["polars>=1.31.0"]
 pandas = ["pandas>=2.2.0"]

vortex-python/pyproject.toml

@@ -0,0 +1,4 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+def launch(args: list[str]) -> None: ...

vortex-python/python/vortex/_lib/cli.pyi

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+import sys
+
+
+def main() -> None:
+    from vortex._lib.cli import launch  # pyright: ignore[reportMissingModuleSource]
+
+    launch(sys.argv)

vortex-python/python/vortex/cli.py

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+
+use crate::SESSION;
+use crate::TOKIO_RUNTIME;
+use crate::install_module;
+
+pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
+    let m = PyModule::new(py, "cli")?;
+    parent.add_submodule(&m)?;
+    install_module("vortex._lib.cli", &m)?;
+
+    m.add_function(wrap_pyfunction!(launch, &m)?)?;
+
+    Ok(())
+}
+
+/// Launch the `vx` CLI with the given arguments.
+///
+/// Parameters
+/// ----------
+/// args : list[str]
+///     Command-line arguments, typically ``sys.argv``.
+#[pyfunction]
+fn launch(py: Python, args: Vec<String>) -> PyResult<()> {
+    py.detach(|| TOKIO_RUNTIME.block_on(vortex_tui::launch_from(&SESSION, args)))
+        .map_err(|e| PyRuntimeError::new_err(e.to_string()))
+}

vortex-python/src/cli.rs

@@ -9,6 +9,7 @@ use pyo3::prelude::*;
 
 pub(crate) mod arrays;
 pub mod arrow;
+mod cli;
 mod compress;
 mod dataset;
 pub(crate) mod dtype;
@@ -61,6 +62,7 @@ fn _lib(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
 
     // Initialize our submodules, living under vortex._lib
     arrays::init(py, m)?;
+    cli::init(py, m)?;
     compress::init(py, m)?;
     dataset::init(py, m)?;
     dtype::init(py, m)?;

vortex-python/src/lib.rs

@@ -22,6 +22,7 @@
 #![deny(clippy::missing_safety_doc)]
 #![deny(missing_docs)]
 
+use std::ffi::OsString;
 use std::path::PathBuf;
 
 use clap::CommandFactory;
@@ -79,13 +80,30 @@ impl Commands {
 
 /// Main entrypoint for `vx` that launches a [`VortexSession`].
 ///
+/// Parses arguments from [`std::env::args_os`]. See [`launch_from`] to supply explicit arguments.
+///
 /// # Errors
 ///
 /// Raises any errors from subcommands.
 pub async fn launch(session: &VortexSession) -> anyhow::Result<()> {
-    env_logger::init();
+    launch_from(session, std::env::args_os()).await
+}
+
+/// Launch `vx` with explicit command-line arguments.
+///
+/// This is useful when embedding the TUI inside another process (e.g. Python) where
+/// [`std::env::args`] may not reflect the intended arguments.
+///
+/// # Errors
+///
+/// Raises any errors from subcommands.
+pub async fn launch_from(
+    session: &VortexSession,
+    args: impl IntoIterator<Item = impl Into<OsString> + Clone>,
+) -> anyhow::Result<()> {
+    let _ = env_logger::try_init();
 
-    let cli = Cli::parse();
+    let cli = Cli::parse_from(args);
 
     let path = cli.command.file_path();
     if !std::fs::exists(path)? {