Allow uv crate to be used as a library

astral-sh · Jun 29, 2024 · 0022fe7 · 0022fe7
1 parent 3a627f3
commit 0022fe7
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 72 deletions.
diff --git a/crates/uv/src/bin/uv.rs b/crates/uv/src/bin/uv.rs
@@ -0,0 +1,7 @@
+use std::process::ExitCode;
+
+use uv::main as uv_main;
+
+fn main() -> ExitCode {
+    uv_main(std::env::args_os())
+}
diff --git a/crates/uv/src/main.rs → crates/uv/src/lib.rs b/crates/uv/src/main.rs → crates/uv/src/lib.rs
@@ -1,4 +1,5 @@
 use std::env;
+use std::ffi::OsString;
 use std::fmt::Write;
 use std::io::stdout;
 use std::path::PathBuf;
@@ -47,76 +48,15 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 #[global_allocator]
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-mod commands;
-mod logging;
-mod printer;
-mod settings;
-mod shell;
-mod version;
-
-#[instrument]
-async fn run() -> Result<ExitStatus> {
-    let cli = match Cli::try_parse() {
-        Ok(cli) => cli,
-        Err(mut err) => {
-            if let Some(ContextValue::String(subcommand)) = err.get(ContextKind::InvalidSubcommand)
-            {
-                match subcommand.as_str() {
-                    "compile" | "lock" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip compile".to_string()),
-                        );
-                    }
-                    "sync" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip sync".to_string()),
-                        );
-                    }
-                    "install" | "add" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip install".to_string()),
-                        );
-                    }
-                    "uninstall" | "remove" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip uninstall".to_string()),
-                        );
-                    }
-                    "freeze" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip freeze".to_string()),
-                        );
-                    }
-                    "list" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip list".to_string()),
-                        );
-                    }
-                    "show" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip show".to_string()),
-                        );
-                    }
-                    "tree" => {
-                        err.insert(
-                            ContextKind::SuggestedSubcommand,
-                            ContextValue::String("uv pip tree".to_string()),
-                        );
-                    }
-                    _ => {}
-                }
-            }
-            err.exit()
-        }
-    };
+pub(crate) mod commands;
+pub(crate) mod logging;
+pub(crate) mod printer;
+pub(crate) mod settings;
+pub(crate) mod shell;
+pub(crate) mod version;
 
+#[instrument(skip_all)]
+async fn run(cli: Cli) -> Result<ExitStatus> {
     // enable flag to pick up warnings generated by workspace loading.
     if !cli.global_args.quiet {
         uv_warnings::enable();
@@ -903,7 +843,74 @@ async fn run() -> Result<ExitStatus> {
     }
 }
 
-fn main() -> ExitCode {
+pub fn main<I, T>(args: I) -> ExitCode
+where
+    I: IntoIterator<Item = T>,
+    T: Into<OsString> + Clone,
+{
+    // `std::env::args` is not `Send` so we parse before passing to our runtime
+    // https://github.com/rust-lang/rust/pull/48005
+    let cli = match Cli::try_parse_from(args) {
+        Ok(cli) => cli,
+        Err(mut err) => {
+            if let Some(ContextValue::String(subcommand)) = err.get(ContextKind::InvalidSubcommand)
+            {
+                match subcommand.as_str() {
+                    "compile" | "lock" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip compile".to_string()),
+                        );
+                    }
+                    "sync" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip sync".to_string()),
+                        );
+                    }
+                    "install" | "add" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip install".to_string()),
+                        );
+                    }
+                    "uninstall" | "remove" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip uninstall".to_string()),
+                        );
+                    }
+                    "freeze" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip freeze".to_string()),
+                        );
+                    }
+                    "list" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip list".to_string()),
+                        );
+                    }
+                    "show" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip show".to_string()),
+                        );
+                    }
+                    "tree" => {
+                        err.insert(
+                            ContextKind::SuggestedSubcommand,
+                            ContextValue::String("uv pip tree".to_string()),
+                        );
+                    }
+                    _ => {}
+                }
+            }
+            err.exit()
+        }
+    };
+
     let result = if let Ok(stack_size) = env::var("UV_STACK_SIZE") {
         // Artificially limit the stack size to test for stack overflows. Windows has a default stack size of 1MB,
         // which is lower than the linux and mac default.
@@ -915,7 +922,7 @@ fn main() -> ExitCode {
                 .thread_stack_size(stack_size)
                 .build()
                 .expect("Failed building the Runtime")
-                .block_on(run())
+                .block_on(run(cli))
         };
         std::thread::Builder::new()
             .stack_size(stack_size)
@@ -928,7 +935,7 @@ fn main() -> ExitCode {
             .enable_all()
             .build()
             .expect("Failed building the Runtime")
-            .block_on(run())
+            .block_on(run(cli))
     };
 
     match result {