From 9efbc1fc25f3ccea269366c45ec9b0177c689ce9 Mon Sep 17 00:00:00 2001 From: Inada Naoki Date: Tue, 20 Feb 2024 04:13:20 +0900 Subject: [PATCH] Add support for `venv --prompt` (#1570) ## Summary This PR adds the `--prompt` option to `venv` subcommand. The default behavior for `uv venv` is to create a virtual environment in the current directory with `.venv` name. This is different from `venv` / `virtualenv` where a user always needs to provide the virtual environment path. This allows us to define our own behavior in the default scenario (`uv venv`). We've decided to use the current directory's name in that case. Workflows: | Command | Virtual Environment Name | Prompt | |--------|--------|--------| | `uv venv` | `.venv` (default) | Current directory name | | `uv venv project` | `project` | `project` | | `uv venv --prompt .` | `.venv` | Current directory name | | `uv venv --prompt foobar` | `.venv` | `foobar` | | `uv venv project --prompt foobar` | `project` | `foobar` | Fixes #1445 ## Test Plan This is my first Rust code and I don't know how to write tests yet. I just checked the behavior manually: ``` $ cargo build $ mkdir t $ cd t $ ../target/debug/uv venv -p 3.11 $ rg -w t .venv/bin/acti* .venv/bin/activate.csh 13:setenv VIRTUAL_ENV '/Users/inada-n/work/uv/t/.venv' 20:if ('t' != "") then 21: setenv VIRTUAL_ENV_PROMPT 't' 23: setenv VIRTUAL_ENV_PROMPT "$VIRTUAL_ENV:t:q" 38: # in which case, $prompt is undefined and we wouldn't .venv/bin/activate 48:VIRTUAL_ENV='/Users/inada-n/work/uv/t/.venv' 59: VIRTUAL_ENV_PROMPT="t" .venv/bin/activate.fish 61:set -gx VIRTUAL_ENV '/Users/inada-n/work/uv/t/.venv' 73:if test -n 't' 74: set -gx VIRTUAL_ENV_PROMPT 't' .venv/bin/activate.ps1 40:if ("t" -ne "") { 41: $env:VIRTUAL_ENV_PROMPT = "t" .venv/bin/activate.nu 6:# but then simply `deactivate` won't work because it is just an alias to hide 35: let virtual_env = '/Users/inada-n/work/uv/t/.venv' 50: let virtual_env_prompt = (if ('t' | is-empty) { 53: 't' ``` --------- Co-authored-by: Dhruv Manilawala --- crates/gourgeist/src/activator/activate | 4 +-- crates/gourgeist/src/activator/activate.bat | 4 +-- crates/gourgeist/src/activator/activate.csh | 4 +-- crates/gourgeist/src/activator/activate.fish | 4 +-- crates/gourgeist/src/activator/activate.nu | 4 +-- crates/gourgeist/src/activator/activate.ps1 | 4 +-- crates/gourgeist/src/bare.rs | 28 ++++++++++++++++-- crates/gourgeist/src/lib.rs | 31 ++++++++++++++++++-- crates/gourgeist/src/main.rs | 6 ++-- crates/uv-build/src/lib.rs | 6 +++- crates/uv/src/commands/venv.rs | 6 +++- crates/uv/src/main.rs | 27 ++++++++++++++++- 12 files changed, 106 insertions(+), 22 deletions(-) diff --git a/crates/gourgeist/src/activator/activate b/crates/gourgeist/src/activator/activate index 237973829..b026a2157 100644 --- a/crates/gourgeist/src/activator/activate +++ b/crates/gourgeist/src/activator/activate @@ -76,8 +76,8 @@ _OLD_VIRTUAL_PATH="$PATH" PATH="$VIRTUAL_ENV/{{ BIN_NAME }}:$PATH" export PATH -if [ "x" != x ] ; then - VIRTUAL_ENV_PROMPT="" +if [ "x{{ VIRTUAL_PROMPT }}" != x ] ; then + VIRTUAL_ENV_PROMPT="{{ VIRTUAL_PROMPT }}" else VIRTUAL_ENV_PROMPT=$(basename "$VIRTUAL_ENV") fi diff --git a/crates/gourgeist/src/activator/activate.bat b/crates/gourgeist/src/activator/activate.bat index 524f61d5a..584627c95 100644 --- a/crates/gourgeist/src/activator/activate.bat +++ b/crates/gourgeist/src/activator/activate.bat @@ -21,7 +21,7 @@ @set "VIRTUAL_ENV={{ VIRTUAL_ENV_DIR }}" -@set "VIRTUAL_ENV_PROMPT=venv" +@set "VIRTUAL_ENV_PROMPT={{ VIRTUAL_PROMPT }}" @if NOT DEFINED VIRTUAL_ENV_PROMPT ( @for %%d in ("%VIRTUAL_ENV%") do @set "VIRTUAL_ENV_PROMPT=%%~nxd" ) @@ -56,4 +56,4 @@ @set "_OLD_VIRTUAL_PATH=%PATH%" :ENDIFVPATH2 -@set "PATH=%VIRTUAL_ENV%\{{ BIN_NAME }};%PATH%" \ No newline at end of file +@set "PATH=%VIRTUAL_ENV%\{{ BIN_NAME }};%PATH%" diff --git a/crates/gourgeist/src/activator/activate.csh b/crates/gourgeist/src/activator/activate.csh index 2453e31eb..24c591def 100644 --- a/crates/gourgeist/src/activator/activate.csh +++ b/crates/gourgeist/src/activator/activate.csh @@ -38,8 +38,8 @@ setenv PATH "$VIRTUAL_ENV:q/{{ BIN_NAME }}:$PATH:q" -if ('' != "") then - setenv VIRTUAL_ENV_PROMPT '' +if ('{{ VIRTUAL_PROMPT }}' != "") then + setenv VIRTUAL_ENV_PROMPT '{{ VIRTUAL_PROMPT }}' else setenv VIRTUAL_ENV_PROMPT "$VIRTUAL_ENV:t:q" endif diff --git a/crates/gourgeist/src/activator/activate.fish b/crates/gourgeist/src/activator/activate.fish index 54a39e0e2..b2202c2d6 100644 --- a/crates/gourgeist/src/activator/activate.fish +++ b/crates/gourgeist/src/activator/activate.fish @@ -91,8 +91,8 @@ set -gx PATH "$VIRTUAL_ENV"'/{{ BIN_NAME }}' $PATH # Prompt override provided? # If not, just use the environment name. -if test -n '' - set -gx VIRTUAL_ENV_PROMPT '' +if test -n '{{ VIRTUAL_PROMPT }}' + set -gx VIRTUAL_ENV_PROMPT '{{ VIRTUAL_PROMPT }}' else set -gx VIRTUAL_ENV_PROMPT (basename "$VIRTUAL_ENV") end diff --git a/crates/gourgeist/src/activator/activate.nu b/crates/gourgeist/src/activator/activate.nu index 61e7e2466..414d86698 100644 --- a/crates/gourgeist/src/activator/activate.nu +++ b/crates/gourgeist/src/activator/activate.nu @@ -68,10 +68,10 @@ export-env { let new_path = ($env | get $path_name | prepend $venv_path) # If there is no default prompt, then use the env name instead - let virtual_env_prompt = (if ('' | is-empty) { + let virtual_env_prompt = (if ('{{ VIRTUAL_PROMPT }}' | is-empty) { ($virtual_env | path basename) } else { - '' + '{{ VIRTUAL_PROMPT }}' }) let new_env = { diff --git a/crates/gourgeist/src/activator/activate.ps1 b/crates/gourgeist/src/activator/activate.ps1 index 6e22a9ebc..ba2213865 100644 --- a/crates/gourgeist/src/activator/activate.ps1 +++ b/crates/gourgeist/src/activator/activate.ps1 @@ -58,8 +58,8 @@ deactivate -nondestructive $VIRTUAL_ENV = $BASE_DIR $env:VIRTUAL_ENV = $VIRTUAL_ENV -if ("" -ne "") { - $env:VIRTUAL_ENV_PROMPT = "" +if ("{{ VIRTUAL_PROMPT }}" -ne "") { + $env:VIRTUAL_ENV_PROMPT = "{{ VIRTUAL_PROMPT }}" } else { $env:VIRTUAL_ENV_PROMPT = $( Split-Path $env:VIRTUAL_ENV -Leaf ) diff --git a/crates/gourgeist/src/bare.rs b/crates/gourgeist/src/bare.rs index 91831af95..3914d58ca 100644 --- a/crates/gourgeist/src/bare.rs +++ b/crates/gourgeist/src/bare.rs @@ -1,5 +1,6 @@ //! Create a bare virtualenv without any packages install +use std::env; use std::env::consts::EXE_SUFFIX; use std::io; use std::io::{BufWriter, Write}; @@ -11,6 +12,8 @@ use tracing::info; use uv_interpreter::Interpreter; +use crate::Prompt; + /// The bash activate scripts with the venv dependent paths patches out const ACTIVATE_TEMPLATES: &[(&str, &str)] = &[ ("activate", include_str!("activator/activate")), @@ -29,10 +32,17 @@ const ACTIVATE_TEMPLATES: &[(&str, &str)] = &[ const VIRTUALENV_PATCH: &str = include_str!("_virtualenv.py"); /// Very basic `.cfg` file format writer. -fn write_cfg(f: &mut impl Write, data: &[(&str, String); 8]) -> io::Result<()> { +fn write_cfg( + f: &mut impl Write, + data: &[(&str, String); 8], + prompt: Option, +) -> io::Result<()> { for (key, value) in data { writeln!(f, "{key} = {value}")?; } + if let Some(prompt) = prompt { + writeln!(f, "prompt = {prompt}")?; + } Ok(()) } @@ -58,7 +68,11 @@ pub struct VenvPaths { } /// Write all the files that belong to a venv without any packages installed. -pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::Result { +pub fn create_bare_venv( + location: &Utf8Path, + interpreter: &Interpreter, + prompt: Prompt, +) -> io::Result { // We have to canonicalize the interpreter path, otherwise the home is set to the venv dir instead of the real root. // This would make python-build-standalone fail with the encodings module not being found because its home is wrong. let base_python: Utf8PathBuf = fs_err::canonicalize(interpreter.sys_executable())? @@ -107,6 +121,13 @@ pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::R unimplemented!("Only Windows and Unix are supported") }; let bin_dir = location.join(bin_name); + let prompt = match prompt { + Prompt::CurrentDirectoryName => env::current_dir()? + .file_name() + .map(|name| name.to_string_lossy().to_string()), + Prompt::Static(value) => Some(value), + Prompt::None => None, + }; // Add the CACHEDIR.TAG. cachedir::ensure_tag(&location)?; @@ -164,6 +185,7 @@ pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::R let activator = template .replace("{{ VIRTUAL_ENV_DIR }}", location.as_str()) .replace("{{ BIN_NAME }}", bin_name) + .replace("{{ VIRTUAL_PROMPT }}", prompt.as_deref().unwrap_or("")) .replace( "{{ RELATIVE_SITE_PACKAGES }}", &format!( @@ -217,7 +239,7 @@ pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::R ("base-executable", base_python.to_string()), ]; let mut pyvenv_cfg = BufWriter::new(File::create(location.join("pyvenv.cfg"))?); - write_cfg(&mut pyvenv_cfg, pyvenv_cfg_data)?; + write_cfg(&mut pyvenv_cfg, pyvenv_cfg_data, prompt)?; drop(pyvenv_cfg); let site_packages = if cfg!(unix) { diff --git a/crates/gourgeist/src/lib.rs b/crates/gourgeist/src/lib.rs index 3a41c8686..954525302 100644 --- a/crates/gourgeist/src/lib.rs +++ b/crates/gourgeist/src/lib.rs @@ -23,12 +23,39 @@ pub enum Error { Platform(#[from] PlatformError), } +/// The value to use for the shell prompt when inside a virtual environment. +#[derive(Debug)] +pub enum Prompt { + /// Use the current directory name as the prompt. + CurrentDirectoryName, + /// Use the fixed string as the prompt. + Static(String), + /// Default to no prompt. The prompt is then set by the activator script + /// to the virtual environment's directory name. + None, +} + +impl Prompt { + /// Determine the prompt value to be used from the command line arguments. + pub fn from_args(prompt: Option) -> Self { + match prompt { + Some(prompt) if prompt == "." => Prompt::CurrentDirectoryName, + Some(prompt) => Prompt::Static(prompt), + None => Prompt::None, + } + } +} + /// Create a virtualenv. -pub fn create_venv(location: &Path, interpreter: Interpreter) -> Result { +pub fn create_venv( + location: &Path, + interpreter: Interpreter, + prompt: Prompt, +) -> Result { let location: &Utf8Path = location .try_into() .map_err(|err: FromPathError| err.into_io_error())?; - let paths = create_bare_venv(location, &interpreter)?; + let paths = create_bare_venv(location, &interpreter, prompt)?; Ok(Virtualenv::from_interpreter( interpreter, paths.root.as_std_path(), diff --git a/crates/gourgeist/src/main.rs b/crates/gourgeist/src/main.rs index cb1e90b06..a6718fef2 100644 --- a/crates/gourgeist/src/main.rs +++ b/crates/gourgeist/src/main.rs @@ -11,7 +11,7 @@ use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; use tracing_subscriber::{fmt, EnvFilter}; -use gourgeist::{create_bare_venv, parse_python_cli}; +use gourgeist::{create_bare_venv, parse_python_cli, Prompt}; use platform_host::Platform; use uv_cache::Cache; use uv_interpreter::Interpreter; @@ -21,6 +21,8 @@ struct Cli { path: Option, #[clap(short, long)] python: Option, + #[clap(long)] + prompt: Option, } fn run() -> Result<(), gourgeist::Error> { @@ -34,7 +36,7 @@ fn run() -> Result<(), gourgeist::Error> { Cache::from_path(".gourgeist_cache")? }; let info = Interpreter::query(python.as_std_path(), &platform, &cache).unwrap(); - create_bare_venv(&location, &info)?; + create_bare_venv(&location, &info, Prompt::from_args(cli.prompt))?; Ok(()) } diff --git a/crates/uv-build/src/lib.rs b/crates/uv-build/src/lib.rs index 6539886dc..a4527839f 100644 --- a/crates/uv-build/src/lib.rs +++ b/crates/uv-build/src/lib.rs @@ -323,7 +323,11 @@ impl SourceBuild { let pep517_backend = Self::get_pep517_backend(setup_py, &source_tree, &default_backend) .map_err(|err| *err)?; - let venv = gourgeist::create_venv(&temp_dir.path().join(".venv"), interpreter.clone())?; + let venv = gourgeist::create_venv( + &temp_dir.path().join(".venv"), + interpreter.clone(), + gourgeist::Prompt::None, + )?; // Setup the build environment. let resolved_requirements = Self::get_resolved_requirements( diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs index 1111e6d88..43b7214aa 100644 --- a/crates/uv/src/commands/venv.rs +++ b/crates/uv/src/commands/venv.rs @@ -11,6 +11,7 @@ use owo_colors::OwoColorize; use thiserror::Error; use distribution_types::{DistributionMetadata, IndexLocations, Name}; +use gourgeist::Prompt; use pep508_rs::Requirement; use platform_host::Platform; use uv_cache::Cache; @@ -31,6 +32,7 @@ pub(crate) async fn venv( path: &Path, python_request: Option<&str>, index_locations: &IndexLocations, + prompt: Prompt, connectivity: Connectivity, seed: bool, exclude_newer: Option>, @@ -41,6 +43,7 @@ pub(crate) async fn venv( path, python_request, index_locations, + prompt, connectivity, seed, exclude_newer, @@ -82,6 +85,7 @@ async fn venv_impl( path: &Path, python_request: Option<&str>, index_locations: &IndexLocations, + prompt: Prompt, connectivity: Connectivity, seed: bool, exclude_newer: Option>, @@ -115,7 +119,7 @@ async fn venv_impl( .into_diagnostic()?; // Create the virtual environment. - let venv = gourgeist::create_venv(path, interpreter).map_err(VenvError::Creation)?; + let venv = gourgeist::create_venv(path, interpreter, prompt).map_err(VenvError::Creation)?; // Install seed packages. if seed { diff --git a/crates/uv/src/main.rs b/crates/uv/src/main.rs index b296c2287..95e3a693f 100644 --- a/crates/uv/src/main.rs +++ b/crates/uv/src/main.rs @@ -49,6 +49,8 @@ mod logging; mod printer; mod requirements; +const DEFAULT_VENV_NAME: &str = ".venv"; + #[derive(Parser)] #[command(author, version, about)] #[command(propagate_version = true)] @@ -652,9 +654,21 @@ struct VenvArgs { seed: bool, /// The path to the virtual environment to create. - #[clap(default_value = ".venv")] + #[clap(default_value = DEFAULT_VENV_NAME)] name: PathBuf, + /// Provide an alternative prompt prefix for the virtual environment. + /// + /// The default behavior depends on whether the virtual environment path is provided: + /// - If provided (`uv venv project`), the prompt is set to the virtual environment's directory name. + /// - If not provided (`uv venv`), the prompt is set to the current directory's name. + /// + /// Possible values: + /// - `.`: Use the current directory name. + /// - Any string: Use the given string. + #[clap(long, verbatim_doc_comment)] + prompt: Option, + /// The URL of the Python Package Index. #[clap(long, short, default_value = IndexUrl::Pypi.as_str(), env = "UV_INDEX_URL")] index_url: IndexUrl, @@ -1016,10 +1030,21 @@ async fn run() -> Result { Vec::new(), args.no_index, ); + + // Since we use ".venv" as the default name, we use "." as the default prompt. + let prompt = args.prompt.or_else(|| { + if args.name == PathBuf::from(DEFAULT_VENV_NAME) { + Some(".".to_string()) + } else { + None + } + }); + commands::venv( &args.name, args.python.as_deref(), &index_locations, + gourgeist::Prompt::from_args(prompt), if args.offline { Connectivity::Offline } else {