Skip to content

Commit b5b7cc8

Browse files
committed
feat: add mamba/micromamba manager detection and reporting (Fixes #25)
- Add Mamba variant to EnvManagerType enum - Detect mamba/micromamba binaries alongside conda in install directories - Report both conda and mamba managers when both exist - Detect when conda_executable points to mamba for backwards compatibility - Discover mamba/micromamba on PATH for environment location discovery - Fall back to mamba when conda binary unavailable for missing env detection - Update JSONRPC docs with Mamba manager type
1 parent 9202a5d commit b5b7cc8

File tree

6 files changed

+143
-10
lines changed

6 files changed

+143
-10
lines changed

crates/pet-conda/src/environment_locations.rs

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
use crate::{
55
conda_rc::{get_conda_rc_search_paths, Condarc},
66
env_variables::EnvVariables,
7-
manager::find_conda_binary,
7+
manager::{find_conda_binary, find_mamba_binary},
88
utils::{is_conda_env, is_conda_install},
99
};
1010
use log::trace;
@@ -372,6 +372,10 @@ pub fn get_known_conda_install_locations(
372372
if let Some(conda_dir) = get_conda_dir_from_exe(&conda_from_path) {
373373
known_paths.push(conda_dir);
374374
}
375+
// Also check for mamba/micromamba on PATH to discover its install directory
376+
if let Some(mamba_dir) = get_conda_dir_from_exe(&find_mamba_binary(env_vars)) {
377+
known_paths.push(mamba_dir);
378+
}
375379
known_paths.sort();
376380
known_paths.dedup();
377381

@@ -456,6 +460,10 @@ pub fn get_known_conda_install_locations(
456460
if let Some(conda_dir) = get_conda_dir_from_exe(&conda_from_path) {
457461
known_paths.push(conda_dir);
458462
}
463+
// Also check for mamba/micromamba on PATH to discover its install directory
464+
if let Some(mamba_dir) = get_conda_dir_from_exe(&find_mamba_binary(env_vars)) {
465+
known_paths.push(mamba_dir);
466+
}
459467
known_paths.sort();
460468
known_paths.dedup();
461469
known_paths.into_iter().filter(|f| f.exists()).collect()

crates/pet-conda/src/lib.rs

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use environment_locations::{
99
};
1010
use environments::{get_conda_environment_info, CondaEnvironment};
1111
use log::error;
12-
use manager::CondaManager;
12+
use manager::{get_mamba_manager, is_mamba_executable, CondaManager};
1313
use pet_core::{
1414
cache::LocatorCache,
1515
env::PythonEnv,
@@ -64,6 +64,7 @@ pub struct CondaTelemetryInfo {
6464
pub struct Conda {
6565
pub environments: Arc<LocatorCache<PathBuf, PythonEnvironment>>,
6666
pub managers: Arc<LocatorCache<PathBuf, CondaManager>>,
67+
pub mamba_managers: Arc<LocatorCache<PathBuf, CondaManager>>,
6768
pub env_vars: EnvVariables,
6869
conda_executable: Arc<RwLock<Option<PathBuf>>>,
6970
}
@@ -73,13 +74,15 @@ impl Conda {
7374
Conda {
7475
environments: Arc::new(LocatorCache::new()),
7576
managers: Arc::new(LocatorCache::new()),
77+
mamba_managers: Arc::new(LocatorCache::new()),
7678
env_vars: EnvVariables::from(env),
7779
conda_executable: Arc::new(RwLock::new(None)),
7880
}
7981
}
8082
fn clear(&self) {
8183
self.environments.clear();
8284
self.managers.clear();
85+
self.mamba_managers.clear();
8386
}
8487
}
8588

@@ -91,7 +94,12 @@ impl CondaLocator for Conda {
9194
) -> Option<()> {
9295
// Look for environments that we couldn't find without spawning conda.
9396
let user_provided_conda_exe = conda_executable.is_some();
94-
let conda_info = CondaInfo::from(conda_executable)?;
97+
// Try the provided executable first (could be conda or mamba for backwards compat),
98+
// then fall back to mamba/micromamba found on PATH if conda is unavailable.
99+
let conda_info = CondaInfo::from(conda_executable).or_else(|| {
100+
let mamba_exe = manager::find_mamba_binary(&self.env_vars);
101+
CondaInfo::from(mamba_exe)
102+
})?;
95103
let environments_map = self.environments.clone_map();
96104
let new_envs = conda_info
97105
.envs
@@ -159,6 +167,15 @@ impl CondaLocator for Conda {
159167
// Possible we'll find environments in other directories created using this manager
160168
self.managers.insert(conda_dir.clone(), manager.clone());
161169

170+
// Also check for a mamba/micromamba manager in the same directory and report it.
171+
if !self.mamba_managers.contains_key(&conda_dir) {
172+
if let Some(mamba_mgr) = get_mamba_manager(&conda_dir) {
173+
self.mamba_managers
174+
.insert(conda_dir.clone(), mamba_mgr.clone());
175+
reporter.report_manager(&mamba_mgr.to_manager());
176+
}
177+
}
178+
162179
// Find all the environments in the conda install folder. (under `envs` folder)
163180
for conda_env in
164181
get_conda_environments(&get_environments(&conda_dir), &manager.clone().into())
@@ -272,6 +289,18 @@ impl Locator for Conda {
272289
let env_vars = self.env_vars.clone();
273290
let executable = self.conda_executable.read().unwrap().clone();
274291
thread::scope(|s| {
292+
// If the user-provided conda_executable is actually a mamba/micromamba binary
293+
// (backwards compatibility), report it as a mamba manager and discover its envs.
294+
if let Some(ref exe) = executable {
295+
if is_mamba_executable(exe) {
296+
if let Some(mamba_dir) = get_conda_dir_from_exe(&executable) {
297+
if let Some(mamba_mgr) = get_mamba_manager(&mamba_dir) {
298+
self.mamba_managers.insert(mamba_dir, mamba_mgr.clone());
299+
reporter.report_manager(&mamba_mgr.to_manager());
300+
}
301+
}
302+
}
303+
}
275304
// 1. Get a list of all know conda environments file paths
276305
let possible_conda_envs = get_conda_environment_paths(&env_vars, &executable);
277306
for path in possible_conda_envs {
@@ -318,6 +347,13 @@ impl Locator for Conda {
318347
self.environments.insert(prefix.clone(), env.clone());
319348
reporter.report_manager(&manager.to_manager());
320349
reporter.report_environment(&env);
350+
351+
// Also check for a mamba/micromamba manager in the same directory and report it.
352+
if let Some(mamba_mgr) = self.mamba_managers.get_or_insert_with(conda_dir.clone(), || {
353+
get_mamba_manager(conda_dir)
354+
}) {
355+
reporter.report_manager(&mamba_mgr.to_manager());
356+
}
321357
} else {
322358
// We will still return the conda env even though we do not have the manager.
323359
// This might seem incorrect, however the tool is about discovering environments.

crates/pet-conda/src/manager.rs

Lines changed: 88 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,30 @@ fn get_conda_executable(path: &Path) -> Option<PathBuf> {
3636
None
3737
}
3838

39+
fn get_mamba_executable(path: &Path) -> Option<PathBuf> {
40+
#[cfg(windows)]
41+
let relative_paths = vec![
42+
PathBuf::from("Scripts").join("mamba.exe"),
43+
PathBuf::from("Scripts").join("micromamba.exe"),
44+
PathBuf::from("bin").join("mamba.exe"),
45+
PathBuf::from("bin").join("micromamba.exe"),
46+
];
47+
#[cfg(unix)]
48+
let relative_paths = vec![
49+
PathBuf::from("bin").join("mamba"),
50+
PathBuf::from("bin").join("micromamba"),
51+
];
52+
53+
for relative_path in relative_paths {
54+
let exe = path.join(&relative_path);
55+
if exe.exists() {
56+
return Some(exe);
57+
}
58+
}
59+
60+
None
61+
}
62+
3963
/// Specifically returns the file names that are valid for 'conda' on windows
4064
#[cfg(windows)]
4165
fn get_conda_bin_names() -> Vec<&'static str> {
@@ -48,6 +72,18 @@ fn get_conda_bin_names() -> Vec<&'static str> {
4872
vec!["conda"]
4973
}
5074

75+
/// Specifically returns the file names that are valid for 'mamba'/'micromamba' on windows
76+
#[cfg(windows)]
77+
fn get_mamba_bin_names() -> Vec<&'static str> {
78+
vec!["mamba.exe", "micromamba.exe"]
79+
}
80+
81+
/// Specifically returns the file names that are valid for 'mamba'/'micromamba' on linux/Mac
82+
#[cfg(unix)]
83+
fn get_mamba_bin_names() -> Vec<&'static str> {
84+
vec!["mamba", "micromamba"]
85+
}
86+
5187
/// Find the conda binary on the PATH environment variable
5288
pub fn find_conda_binary(env_vars: &EnvVariables) -> Option<PathBuf> {
5389
let paths = env_vars.path.clone()?;
@@ -62,17 +98,32 @@ pub fn find_conda_binary(env_vars: &EnvVariables) -> Option<PathBuf> {
6298
None
6399
}
64100

101+
/// Find a mamba or micromamba binary on the PATH environment variable
102+
pub fn find_mamba_binary(env_vars: &EnvVariables) -> Option<PathBuf> {
103+
let paths = env_vars.path.clone()?;
104+
for path in env::split_paths(&paths) {
105+
for bin in get_mamba_bin_names() {
106+
let mamba_path = path.join(bin);
107+
if mamba_path.is_file() || mamba_path.is_symlink() {
108+
return Some(mamba_path);
109+
}
110+
}
111+
}
112+
None
113+
}
114+
65115
#[derive(Debug, Clone)]
66116
pub struct CondaManager {
67117
pub executable: PathBuf,
68118
pub version: Option<String>,
69119
pub conda_dir: Option<PathBuf>,
120+
pub manager_type: EnvManagerType,
70121
}
71122

72123
impl CondaManager {
73124
pub fn to_manager(&self) -> EnvManager {
74125
EnvManager {
75-
tool: EnvManagerType::Conda,
126+
tool: self.manager_type,
76127
executable: self.executable.clone(),
77128
version: self.version.clone(),
78129
}
@@ -85,7 +136,9 @@ impl CondaManager {
85136
// If this environment is in a folder named `envs`, then the parent directory of `envs` is the root conda install folder.
86137
if let Some(parent) = path.ancestors().nth(2) {
87138
if is_conda_install(parent) {
88-
if let Some(manager) = get_conda_manager(parent) {
139+
if let Some(manager) =
140+
get_conda_manager(parent).or_else(|| get_mamba_manager(parent))
141+
{
89142
return Some(manager);
90143
}
91144
}
@@ -98,6 +151,7 @@ impl CondaManager {
98151
// Or its in a location such as `~/.conda/envs` or `~/miniconda3/envs` where the conda install folder is not a parent of this path.
99152
if let Some(conda_install_folder) = get_conda_installation_used_to_create_conda_env(path) {
100153
get_conda_manager(&conda_install_folder)
154+
.or_else(|| get_mamba_manager(&conda_install_folder))
101155
} else {
102156
// If this is a conda env and the parent is `.conda/envs`, then this is definitely NOT a root conda install folder.
103157
// Hence never use conda installs from these env paths.
@@ -111,19 +165,24 @@ impl CondaManager {
111165
}
112166
}
113167

114-
if let Some(manager) = get_conda_manager(path) {
168+
if let Some(manager) = get_conda_manager(path).or_else(|| get_mamba_manager(path)) {
115169
Some(manager)
116170
} else {
117-
trace!("No conda manager found for path: {:?}", path);
171+
trace!("No conda or mamba manager found for path: {:?}", path);
118172
None
119173
}
120174
}
121175
}
122-
pub fn from_info(executable: &Path, info: &CondaInfo) -> Option<CondaManager> {
176+
pub fn from_info(
177+
executable: &Path,
178+
info: &CondaInfo,
179+
manager_type: EnvManagerType,
180+
) -> Option<CondaManager> {
123181
Some(CondaManager {
124182
executable: executable.to_path_buf(),
125183
version: Some(info.conda_version.clone()),
126184
conda_dir: info.conda_prefix.clone(),
185+
manager_type,
127186
})
128187
}
129188
}
@@ -135,8 +194,32 @@ fn get_conda_manager(path: &Path) -> Option<CondaManager> {
135194
executable: conda_exe,
136195
version: Some(conda_pkg.version),
137196
conda_dir: Some(path.to_path_buf()),
197+
manager_type: EnvManagerType::Conda,
138198
})
139199
} else {
140200
None
141201
}
142202
}
203+
204+
/// Checks whether a given executable path refers to a mamba or micromamba binary.
205+
pub fn is_mamba_executable(exe: &Path) -> bool {
206+
if let Some(name) = exe.file_name().and_then(|n| n.to_str()) {
207+
let name = name.to_lowercase();
208+
name.starts_with("mamba") || name.starts_with("micromamba")
209+
} else {
210+
false
211+
}
212+
}
213+
214+
pub(crate) fn get_mamba_manager(path: &Path) -> Option<CondaManager> {
215+
let mamba_exe = get_mamba_executable(path)?;
216+
// We cannot reliably determine the mamba/micromamba version from package metadata alone.
217+
// The conda package version in conda-meta is the conda version, not the mamba version.
218+
// Determining the mamba version would require spawning the mamba process.
219+
Some(CondaManager {
220+
executable: mamba_exe,
221+
version: None,
222+
conda_dir: Some(path.to_path_buf()),
223+
manager_type: EnvManagerType::Mamba,
224+
})
225+
}

crates/pet-conda/src/telemetry.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,11 @@ fn log_and_find_missing_envs(
221221
.collect::<Vec<_>>();
222222

223223
// Oh oh, we have new envs, lets see what they are.
224-
let manager = CondaManager::from_info(&conda_info.executable, conda_info)?;
224+
let manager = CondaManager::from_info(
225+
&conda_info.executable,
226+
conda_info,
227+
pet_core::manager::EnvManagerType::Conda,
228+
)?;
225229
for path in missing_envs
226230
.clone()
227231
.iter()

crates/pet-core/src/manager.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use std::path::PathBuf;
77
#[derive(Serialize, Deserialize, Copy, Clone, PartialEq, Eq, Debug, Hash)]
88
pub enum EnvManagerType {
99
Conda,
10+
Mamba,
1011
Pipenv,
1112
Poetry,
1213
Pyenv,

docs/JSONRPC.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ interface ConfigureParams {
5454
* This is the path to the conda executable.
5555
*
5656
* Useful for VS Code so users can configure where they have installed Conda.
57+
* For backwards compatibility, this can also be a path to a mamba or micromamba executable.
5758
*/
5859
condaExecutable?: string;
5960
/**
@@ -270,7 +271,7 @@ interface Manager {
270271
/**
271272
* The type of the Manager.
272273
*/
273-
tool: "Conda" | "Pipenv" | "Poetry" | "Pyenv";
274+
tool: "Conda" | "Mamba" | "Pipenv" | "Poetry" | "Pyenv";
274275
/**
275276
* The version of the manager/tool.
276277
* In the case of conda, this is the version of conda.

0 commit comments

Comments
 (0)