Skip to content

Commit ce922d6

Browse files
authored
feat: add mamba/micromamba manager detection and reporting (Fixes #25) (#350)
Adds mamba/micromamba manager detection alongside conda, so both managers are reported when present. **Key changes:** - Add `Mamba` variant to `EnvManagerType` enum - Detect mamba/micromamba binaries in conda install directories and report as separate managers - Auto-detect when `condaExecutable` points to a mamba binary (backwards compatibility) - Discover mamba/micromamba on PATH for environment location discovery - Fall back to mamba when conda binary is unavailable for missing env detection - Update JSONRPC docs with `Mamba` manager type Environments remain `PythonEnvironmentKind::Conda` — only the manager reporting changes. Fixes #25
1 parent 9202a5d commit ce922d6

File tree

22 files changed

+420
-12
lines changed

22 files changed

+420
-12
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: 45 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,17 @@ 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+
let _ = self
172+
.mamba_managers
173+
.get_or_insert_with(conda_dir.clone(), || {
174+
let mgr = get_mamba_manager(&conda_dir);
175+
if let Some(ref m) = mgr {
176+
reporter.report_manager(&m.to_manager());
177+
}
178+
mgr
179+
});
180+
162181
// Find all the environments in the conda install folder. (under `envs` folder)
163182
for conda_env in
164183
get_conda_environments(&get_environments(&conda_dir), &manager.clone().into())
@@ -272,6 +291,18 @@ impl Locator for Conda {
272291
let env_vars = self.env_vars.clone();
273292
let executable = self.conda_executable.read().unwrap().clone();
274293
thread::scope(|s| {
294+
// If the user-provided conda_executable is actually a mamba/micromamba binary
295+
// (backwards compatibility), report it as a mamba manager and discover its envs.
296+
if let Some(ref exe) = executable {
297+
if is_mamba_executable(exe) {
298+
if let Some(mamba_dir) = get_conda_dir_from_exe(&executable) {
299+
if let Some(mamba_mgr) = get_mamba_manager(&mamba_dir) {
300+
self.mamba_managers.insert(mamba_dir, mamba_mgr.clone());
301+
reporter.report_manager(&mamba_mgr.to_manager());
302+
}
303+
}
304+
}
305+
}
275306
// 1. Get a list of all know conda environments file paths
276307
let possible_conda_envs = get_conda_environment_paths(&env_vars, &executable);
277308
for path in possible_conda_envs {
@@ -318,6 +349,18 @@ impl Locator for Conda {
318349
self.environments.insert(prefix.clone(), env.clone());
319350
reporter.report_manager(&manager.to_manager());
320351
reporter.report_environment(&env);
352+
353+
// Also check for a mamba/micromamba manager in the same directory and report it.
354+
// Reporting inside the closure minimizes the TOCTOU window compared to a
355+
// separate contains_key check, though concurrent threads may still
356+
// briefly both invoke the closure before the write-lock double-check.
357+
let _ = self.mamba_managers.get_or_insert_with(conda_dir.clone(), || {
358+
let mgr = get_mamba_manager(conda_dir);
359+
if let Some(ref m) = mgr {
360+
reporter.report_manager(&m.to_manager());
361+
}
362+
mgr
363+
});
321364
} else {
322365
// We will still return the conda env even though we do not have the manager.
323366
// This might seem incorrect, however the tool is about discovering environments.

crates/pet-conda/src/manager.rs

Lines changed: 92 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,34 @@ 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("mamba.bat"),
44+
PathBuf::from("Scripts").join("micromamba.exe"),
45+
PathBuf::from("Scripts").join("micromamba.bat"),
46+
PathBuf::from("bin").join("mamba.exe"),
47+
PathBuf::from("bin").join("mamba.bat"),
48+
PathBuf::from("bin").join("micromamba.exe"),
49+
PathBuf::from("bin").join("micromamba.bat"),
50+
];
51+
#[cfg(unix)]
52+
let relative_paths = vec![
53+
PathBuf::from("bin").join("mamba"),
54+
PathBuf::from("bin").join("micromamba"),
55+
];
56+
57+
for relative_path in relative_paths {
58+
let exe = path.join(&relative_path);
59+
if exe.exists() {
60+
return Some(exe);
61+
}
62+
}
63+
64+
None
65+
}
66+
3967
/// Specifically returns the file names that are valid for 'conda' on windows
4068
#[cfg(windows)]
4169
fn get_conda_bin_names() -> Vec<&'static str> {
@@ -48,6 +76,18 @@ fn get_conda_bin_names() -> Vec<&'static str> {
4876
vec!["conda"]
4977
}
5078

79+
/// Specifically returns the file names that are valid for 'mamba'/'micromamba' on windows
80+
#[cfg(windows)]
81+
fn get_mamba_bin_names() -> Vec<&'static str> {
82+
vec!["mamba.exe", "mamba.bat", "micromamba.exe", "micromamba.bat"]
83+
}
84+
85+
/// Specifically returns the file names that are valid for 'mamba'/'micromamba' on linux/Mac
86+
#[cfg(unix)]
87+
fn get_mamba_bin_names() -> Vec<&'static str> {
88+
vec!["mamba", "micromamba"]
89+
}
90+
5191
/// Find the conda binary on the PATH environment variable
5292
pub fn find_conda_binary(env_vars: &EnvVariables) -> Option<PathBuf> {
5393
let paths = env_vars.path.clone()?;
@@ -62,17 +102,32 @@ pub fn find_conda_binary(env_vars: &EnvVariables) -> Option<PathBuf> {
62102
None
63103
}
64104

105+
/// Find a mamba or micromamba binary on the PATH environment variable
106+
pub fn find_mamba_binary(env_vars: &EnvVariables) -> Option<PathBuf> {
107+
let paths = env_vars.path.clone()?;
108+
for path in env::split_paths(&paths) {
109+
for bin in get_mamba_bin_names() {
110+
let mamba_path = path.join(bin);
111+
if mamba_path.is_file() || mamba_path.is_symlink() {
112+
return Some(mamba_path);
113+
}
114+
}
115+
}
116+
None
117+
}
118+
65119
#[derive(Debug, Clone)]
66120
pub struct CondaManager {
67121
pub executable: PathBuf,
68122
pub version: Option<String>,
69123
pub conda_dir: Option<PathBuf>,
124+
pub manager_type: EnvManagerType,
70125
}
71126

72127
impl CondaManager {
73128
pub fn to_manager(&self) -> EnvManager {
74129
EnvManager {
75-
tool: EnvManagerType::Conda,
130+
tool: self.manager_type,
76131
executable: self.executable.clone(),
77132
version: self.version.clone(),
78133
}
@@ -85,7 +140,9 @@ impl CondaManager {
85140
// If this environment is in a folder named `envs`, then the parent directory of `envs` is the root conda install folder.
86141
if let Some(parent) = path.ancestors().nth(2) {
87142
if is_conda_install(parent) {
88-
if let Some(manager) = get_conda_manager(parent) {
143+
if let Some(manager) =
144+
get_conda_manager(parent).or_else(|| get_mamba_manager(parent))
145+
{
89146
return Some(manager);
90147
}
91148
}
@@ -98,6 +155,7 @@ impl CondaManager {
98155
// Or its in a location such as `~/.conda/envs` or `~/miniconda3/envs` where the conda install folder is not a parent of this path.
99156
if let Some(conda_install_folder) = get_conda_installation_used_to_create_conda_env(path) {
100157
get_conda_manager(&conda_install_folder)
158+
.or_else(|| get_mamba_manager(&conda_install_folder))
101159
} else {
102160
// If this is a conda env and the parent is `.conda/envs`, then this is definitely NOT a root conda install folder.
103161
// Hence never use conda installs from these env paths.
@@ -111,19 +169,24 @@ impl CondaManager {
111169
}
112170
}
113171

114-
if let Some(manager) = get_conda_manager(path) {
172+
if let Some(manager) = get_conda_manager(path).or_else(|| get_mamba_manager(path)) {
115173
Some(manager)
116174
} else {
117-
trace!("No conda manager found for path: {:?}", path);
175+
trace!("No conda or mamba manager found for path: {:?}", path);
118176
None
119177
}
120178
}
121179
}
122-
pub fn from_info(executable: &Path, info: &CondaInfo) -> Option<CondaManager> {
180+
pub fn from_info(
181+
executable: &Path,
182+
info: &CondaInfo,
183+
manager_type: EnvManagerType,
184+
) -> Option<CondaManager> {
123185
Some(CondaManager {
124186
executable: executable.to_path_buf(),
125187
version: Some(info.conda_version.clone()),
126188
conda_dir: info.conda_prefix.clone(),
189+
manager_type,
127190
})
128191
}
129192
}
@@ -135,8 +198,32 @@ fn get_conda_manager(path: &Path) -> Option<CondaManager> {
135198
executable: conda_exe,
136199
version: Some(conda_pkg.version),
137200
conda_dir: Some(path.to_path_buf()),
201+
manager_type: EnvManagerType::Conda,
138202
})
139203
} else {
140204
None
141205
}
142206
}
207+
208+
/// Checks whether a given executable path refers to a mamba or micromamba binary.
209+
pub fn is_mamba_executable(exe: &Path) -> bool {
210+
if let Some(name) = exe.file_name().and_then(|n| n.to_str()) {
211+
let name = name.to_lowercase();
212+
name.starts_with("mamba") || name.starts_with("micromamba")
213+
} else {
214+
false
215+
}
216+
}
217+
218+
pub(crate) fn get_mamba_manager(path: &Path) -> Option<CondaManager> {
219+
let mamba_exe = get_mamba_executable(path)?;
220+
// We cannot reliably determine the mamba/micromamba version from package metadata alone.
221+
// The conda package version in conda-meta is the conda version, not the mamba version.
222+
// Determining the mamba version would require spawning the mamba process.
223+
Some(CondaManager {
224+
executable: mamba_exe,
225+
version: None,
226+
conda_dir: Some(path.to_path_buf()),
227+
manager_type: EnvManagerType::Mamba,
228+
})
229+
}

crates/pet-conda/src/telemetry.rs

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,12 @@ use pet_core::{
1111
};
1212

1313
use crate::{
14-
conda_info::CondaInfo, conda_rc::Condarc, env_variables::EnvVariables,
15-
environments::get_conda_environment_info, manager::CondaManager, utils::is_conda_install,
14+
conda_info::CondaInfo,
15+
conda_rc::Condarc,
16+
env_variables::EnvVariables,
17+
environments::get_conda_environment_info,
18+
manager::{is_mamba_executable, CondaManager},
19+
utils::is_conda_install,
1620
};
1721

1822
pub fn report_missing_envs(
@@ -221,7 +225,12 @@ fn log_and_find_missing_envs(
221225
.collect::<Vec<_>>();
222226

223227
// Oh oh, we have new envs, lets see what they are.
224-
let manager = CondaManager::from_info(&conda_info.executable, conda_info)?;
228+
let manager_type = if is_mamba_executable(&conda_info.executable) {
229+
pet_core::manager::EnvManagerType::Mamba
230+
} else {
231+
pet_core::manager::EnvManagerType::Conda
232+
};
233+
let manager = CondaManager::from_info(&conda_info.executable, conda_info, manager_type)?;
225234
for path in missing_envs
226235
.clone()
227236
.iter()

0 commit comments

Comments
 (0)