use ini::Ini; use std::path::Path; use super::{Context, Module, RootModuleConfig}; use crate::configs::python::PythonConfig; use crate::formatter::StringFormatter; use crate::formatter::VersionFormatter; use crate::utils::get_command_string_output; /// Creates a module with the current Python version and, if active, virtual environment. pub fn module<'a>(context: &'a Context) -> Option> { let mut module = context.new_module("python"); let config: PythonConfig = PythonConfig::try_load(module.config); let is_py_project = context .try_begin_scan()? .set_files(&config.detect_files) .set_extensions(&config.detect_extensions) .set_folders(&config.detect_folders) .is_match(); let is_venv = context.get_env("VIRTUAL_ENV").is_some(); if !is_py_project && !is_venv { return None; }; let pyenv_prefix = if config.pyenv_version_name { config.pyenv_prefix } else { "" }; let parsed = StringFormatter::new(config.format).and_then(|formatter| { formatter .map_meta(|var, _| match var { "symbol" => Some(config.symbol), _ => None, }) .map_style(|variable| match variable { "style" => Some(Ok(config.style)), _ => None, }) .map(|variable| match variable { "version" => { if config.pyenv_version_name { return get_pyenv_version(context).map(Ok); } let python_version = get_python_version(context, &config)?; VersionFormatter::format_module_version( module.get_name(), &python_version, config.version_format, ) .map(Ok) } "virtualenv" => { let virtual_env = get_python_virtual_env(context); virtual_env.as_ref().map(|e| Ok(e.trim().to_string())) } "pyenv_prefix" => Some(Ok(pyenv_prefix.to_string())), _ => None, }) .parse(None, Some(context)) }); module.set_segments(match parsed { Ok(segments) => segments, Err(error) => { log::warn!("Error in module `python`:\n{}", error); return None; } }); Some(module) } fn get_pyenv_version(context: &Context) -> Option { let mut version_name = context.get_env("PYENV_VERSION"); if version_name.is_none() { version_name = Some( context .exec_cmd("pyenv", &["version-name"])? .stdout .trim() .to_string(), ) } version_name } fn get_python_version(context: &Context, config: &PythonConfig) -> Option { let version = config .python_binary .0 .iter() .find_map(|binary| context.exec_cmd(binary, &["--version"])) .map(get_command_string_output)?; parse_python_version(&version) } fn parse_python_version(python_version_string: &str) -> Option { let version = python_version_string // split into ["Python", "3.8.6", ...] .split_whitespace() // get down to "3.8.6" .nth(1)?; Some(version.to_string()) } fn get_python_virtual_env(context: &Context) -> Option { context.get_env("VIRTUAL_ENV").and_then(|venv| { get_prompt_from_venv(Path::new(&venv)).or_else(|| { Path::new(&venv) .file_name() .map(|filename| String::from(filename.to_str().unwrap_or(""))) }) }) } fn get_prompt_from_venv(venv_path: &Path) -> Option { Ini::load_from_file(venv_path.join("pyvenv.cfg")) .ok()? .general_section() .get("prompt") .map(|prompt| String::from(prompt.trim_matches(&['(', ')'] as &[_]))) } #[cfg(test)] mod tests { use super::*; use crate::test::ModuleRenderer; use ansi_term::Color; use std::fs::{create_dir_all, File}; use std::io; use std::io::Write; #[test] fn test_parse_python_version() { assert_eq!( parse_python_version("Python 3.7.2"), Some("3.7.2".to_string()) ); } #[test] fn test_parse_python_version_is_malformed() { assert_eq!(parse_python_version("Python 3.7"), Some("3.7".to_string())); } #[test] fn test_parse_python_version_anaconda() { assert_eq!( parse_python_version("Python 3.6.10 :: Anaconda, Inc.",), Some("3.6.10".to_string()) ); } #[test] fn test_parse_python_version_pypy() { assert_eq!( parse_python_version( "\ Python 3.7.9 (7e6e2bb30ac5fbdbd443619cae28c51d5c162a02, Nov 24 2020, 10:03:59) [PyPy 7.3.3-beta0 with GCC 10.2.0]", ), Some("3.7.9".to_string()) ); } #[test] fn folder_without_python_files() -> io::Result<()> { let dir = tempfile::tempdir()?; let actual = ModuleRenderer::new("python").path(dir.path()).collect(); let expected = None; assert_eq!(expected, actual); dir.close() } #[test] fn folder_with_python_version() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join(".python-version"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn folder_with_requirements_txt() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("requirements.txt"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn folder_with_pyproject_toml() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("pyproject.toml"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn folder_with_pipfile() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("Pipfile"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn folder_with_tox() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("tox.ini"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn folder_with_setup_py() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("setup.py"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn folder_with_init_py() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("__init__.py"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn folder_with_py_file() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("main.py"))?.sync_all()?; check_python2_renders(&dir, None); check_python3_renders(&dir, None); check_pyenv_renders(&dir, None); check_multiple_binaries_renders(&dir, None); dir.close() } #[test] fn disabled_scan_for_pyfiles_and_folder_with_ignored_py_file() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("foo.py"))?.sync_all()?; let expected = None; let config = toml::toml! { [python] detect_extensions = [] }; let actual = ModuleRenderer::new("python") .path(dir.path()) .config(config) .collect(); assert_eq!(expected, actual); dir.close() } #[test] fn disabled_scan_for_pyfiles_and_folder_with_setup_py() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("setup.py"))?.sync_all()?; let config = toml::toml! { [python] python_binary = "python2" detect_extensions = [] }; check_python2_renders(&dir, Some(config)); let config_python3 = toml::toml! { [python] python_binary = "python3" detect_extensions = [] }; check_python3_renders(&dir, Some(config_python3)); let config_pyenv = toml::toml! { [python] pyenv_version_name = true pyenv_prefix = "test_pyenv " detect_extensions = [] }; check_pyenv_renders(&dir, Some(config_pyenv)); let config_multi = toml::toml! { [python] python_binary = ["python", "python3"] detect_extensions = [] }; check_multiple_binaries_renders(&dir, Some(config_multi)); dir.close() } #[test] fn with_virtual_env() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("main.py"))?.sync_all()?; let actual = ModuleRenderer::new("python") .path(dir.path()) .env("VIRTUAL_ENV", "/foo/bar/my_venv") .collect(); let expected = Some(format!( "via {}", Color::Yellow.bold().paint("🐍 v3.8.0 (my_venv) ") )); assert_eq!(actual, expected); dir.close() } #[test] fn with_active_venv() -> io::Result<()> { let dir = tempfile::tempdir()?; let actual = ModuleRenderer::new("python") .path(dir.path()) .env("VIRTUAL_ENV", "/foo/bar/my_venv") .collect(); let expected = Some(format!( "via {}", Color::Yellow.bold().paint("🐍 v3.8.0 (my_venv) ") )); assert_eq!(actual, expected); dir.close() } #[test] fn with_active_venv_and_prompt() -> io::Result<()> { let dir = tempfile::tempdir()?; create_dir_all(dir.path().join("my_venv"))?; let mut venv_cfg = File::create(dir.path().join("my_venv").join("pyvenv.cfg"))?; venv_cfg.write_all( br#" home = something prompt = 'foo' "#, )?; venv_cfg.sync_all()?; let actual = ModuleRenderer::new("python") .path(dir.path()) .env("VIRTUAL_ENV", dir.path().join("my_venv").to_str().unwrap()) .collect(); let expected = Some(format!( "via {}", Color::Yellow.bold().paint("🐍 v3.8.0 (foo) ") )); assert_eq!(actual, expected); dir.close() } #[test] fn with_active_venv_and_dirty_prompt() -> io::Result<()> { let dir = tempfile::tempdir()?; create_dir_all(dir.path().join("my_venv"))?; let mut venv_cfg = File::create(dir.path().join("my_venv").join("pyvenv.cfg"))?; venv_cfg.write_all( br#" home = something prompt = '(foo)' "#, )?; venv_cfg.sync_all()?; let actual = ModuleRenderer::new("python") .path(dir.path()) .env("VIRTUAL_ENV", dir.path().join("my_venv").to_str().unwrap()) .collect(); let expected = Some(format!( "via {}", Color::Yellow.bold().paint("🐍 v3.8.0 (foo) ") )); assert_eq!(actual, expected); dir.close() } fn check_python2_renders(dir: &tempfile::TempDir, starship_config: Option) { let config = starship_config.unwrap_or(toml::toml! { [python] python_binary = "python2" }); let actual = ModuleRenderer::new("python") .path(dir.path()) .config(config) .collect(); let expected = Some(format!("via {}", Color::Yellow.bold().paint("🐍 v2.7.17 "))); assert_eq!(expected, actual); } fn check_python3_renders(dir: &tempfile::TempDir, starship_config: Option) { let config = starship_config.unwrap_or(toml::toml! { [python] python_binary = "python3" }); let actual = ModuleRenderer::new("python") .path(dir.path()) .config(config) .collect(); let expected = Some(format!("via {}", Color::Yellow.bold().paint("🐍 v3.8.0 "))); assert_eq!(expected, actual); } fn check_multiple_binaries_renders( dir: &tempfile::TempDir, starship_config: Option, ) { let config = starship_config.unwrap_or(toml::toml! { [python] python_binary = ["python", "python3"] }); let actual = ModuleRenderer::new("python") .path(dir.path()) .config(config) .collect(); let expected = Some(format!("via {}", Color::Yellow.bold().paint("🐍 v3.8.0 "))); assert_eq!(expected, actual); } fn check_pyenv_renders(dir: &tempfile::TempDir, starship_config: Option) { let config = starship_config.unwrap_or(toml::toml! { [python] pyenv_version_name = true pyenv_prefix = "test_pyenv " }); let actual = ModuleRenderer::new("python") .path(dir.path()) .config(config) .collect(); let expected = Some(format!( "via {}", Color::Yellow.bold().paint("🐍 test_pyenv system ") )); assert_eq!(expected, actual); } }