Remove duplicates from reported envs & refactor locators (#23407)

Author: DonJayamanne

native_locator/src/common_python.rs

@@ -2,11 +2,9 @@
 // Licensed under the MIT License.
 
 use crate::known::Environment;
-use crate::locator::Locator;
-use crate::messaging::MessageDispatcher;
+use crate::locator::{Locator, LocatorResult};
 use crate::messaging::PythonEnvironment;
 use crate::utils::{self, PythonEnv};
-use std::collections::HashMap;
 use std::env;
 use std::path::PathBuf;
 
@@ -20,73 +18,61 @@ fn get_env_path(python_executable_path: &PathBuf) -> Option<PathBuf> {
 }
 
 pub struct PythonOnPath<'a> {
-    pub environments: HashMap<String, PythonEnvironment>,
     pub environment: &'a dyn Environment,
 }
 
 impl PythonOnPath<'_> {
     pub fn with<'a>(environment: &'a impl Environment) -> PythonOnPath {
-        PythonOnPath {
-            environments: HashMap::new(),
-            environment,
-        }
+        PythonOnPath { environment }
     }
 }
 
 impl Locator for PythonOnPath<'_> {
-    fn is_known(&self, python_executable: &PathBuf) -> bool {
-        self.environments
-            .contains_key(python_executable.to_str().unwrap_or_default())
-    }
-
-    fn track_if_compatible(&mut self, env: &PythonEnv) -> bool {
+    fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
         let bin = if cfg!(windows) {
             "python.exe"
         } else {
             "python"
         };
         if env.executable.file_name().unwrap().to_ascii_lowercase() != bin {
-            return false;
+            return None;
         }
-        self.environments.insert(
-            env.executable.to_str().unwrap().to_string(),
-            PythonEnvironment {
-                name: None,
-                python_executable_path: Some(env.executable.clone()),
-                version: env.version.clone(),
-                category: crate::messaging::PythonEnvironmentCategory::System,
-                sys_prefix_path: None,
-                env_path: env.path.clone(),
-                env_manager: None,
-                project_path: None,
-                python_run_command: Some(vec![env.executable.to_str().unwrap().to_string()]),
-            },
-        );
-        true
+        Some(PythonEnvironment {
+            name: None,
+            python_executable_path: Some(env.executable.clone()),
+            version: env.version.clone(),
+            category: crate::messaging::PythonEnvironmentCategory::System,
+            sys_prefix_path: None,
+            env_path: env.path.clone(),
+            env_manager: None,
+            project_path: None,
+            python_run_command: Some(vec![env.executable.to_str().unwrap().to_string()]),
+        })
     }
 
-    fn gather(&mut self) -> Option<()> {
+    fn find(&self) -> Option<LocatorResult> {
         let paths = self.environment.get_env_var("PATH".to_string())?;
         let bin = if cfg!(windows) {
             "python.exe"
         } else {
             "python"
         };
+        let mut environments: Vec<PythonEnvironment> = vec![];
         env::split_paths(&paths)
             .map(|p| p.join(bin))
             .filter(|p| p.exists())
             .for_each(|full_path| {
                 let version = utils::get_version(&full_path);
                 let env_path = get_env_path(&full_path);
-                self.track_if_compatible(&PythonEnv::new(full_path, env_path, version));
+                if let Some(env) = self.resolve(&PythonEnv::new(full_path, env_path, version)) {
+                    environments.push(env);
+                }
             });
 
-        Some(())
-    }
-
-    fn report(&self, reporter: &mut dyn MessageDispatcher) {
-        for env in self.environments.values() {
-            reporter.report_environment(env.clone());
+        if environments.is_empty() {
+            None
+        } else {
+            Some(LocatorResult::Environments(environments))
         }
     }
 }

native_locator/src/conda.rs

@@ -4,15 +4,14 @@
 use crate::known;
 use crate::known::Environment;
 use crate::locator::Locator;
+use crate::locator::LocatorResult;
 use crate::messaging;
 use crate::messaging::EnvManager;
 use crate::messaging::EnvManagerType;
-use crate::messaging::MessageDispatcher;
 use crate::messaging::PythonEnvironment;
 use crate::utils::find_python_binary_path;
 use crate::utils::PythonEnv;
 use regex::Regex;
-use std::collections::HashMap;
 use std::env;
 use std::path::{Path, PathBuf};
 
@@ -385,42 +384,35 @@ fn get_distinct_conda_envs(
 }
 
 pub struct Conda<'a> {
-    pub environments: HashMap<String, PythonEnvironment>,
     pub manager: Option<EnvManager>,
     pub environment: &'a dyn Environment,
 }
 
 impl Conda<'_> {
     pub fn with<'a>(environment: &'a impl Environment) -> Conda {
         Conda {
-            environments: HashMap::new(),
             environment,
             manager: None,
         }
     }
 }
 
 impl Locator for Conda<'_> {
-    fn is_known(&self, python_executable: &PathBuf) -> bool {
-        self.environments
-            .contains_key(python_executable.to_str().unwrap_or_default())
-    }
-
-    fn track_if_compatible(&mut self, _env: &PythonEnv) -> bool {
+    fn resolve(&self, _env: &PythonEnv) -> Option<PythonEnvironment> {
         // We will find everything in gather
-        false
+        None
     }
 
-    fn gather(&mut self) -> Option<()> {
+    fn find(&self) -> Option<LocatorResult> {
         let conda_binary = find_conda_binary(self.environment)?;
         let manager = EnvManager::new(
             conda_binary.clone(),
             get_conda_version(&conda_binary),
             EnvManagerType::Conda,
         );
-        self.manager = Some(manager.clone());
 
         let envs = get_distinct_conda_envs(&conda_binary, self.environment);
+        let mut environments: Vec<PythonEnvironment> = vec![];
         for env in envs {
             let executable = find_python_binary_path(Path::new(&env.path));
             let env = messaging::PythonEnvironment::new(
@@ -450,25 +442,13 @@ impl Locator for Conda<'_> {
                 },
             );
 
-            if let Some(exe) = executable {
-                self.environments
-                    .insert(exe.to_str().unwrap_or_default().to_string(), env);
-            } else if let Some(env_path) = env.env_path.clone() {
-                self.environments
-                    .insert(env_path.to_str().unwrap().to_string(), env);
-            }
-        }
-
-        Some(())
-    }
-
-    fn report(&self, reporter: &mut dyn MessageDispatcher) {
-        if let Some(manager) = &self.manager {
-            reporter.report_environment_manager(manager.clone());
+            environments.push(env)
         }
 
-        for env in self.environments.values() {
-            reporter.report_environment(env.clone());
+        if environments.is_empty() {
+            Some(LocatorResult::Managers(vec![manager]))
+        } else {
+            Some(LocatorResult::Environments(environments))
         }
     }
 }

native_locator/src/homebrew.rs

@@ -1,20 +1,14 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use std::{
-    collections::{HashMap, HashSet},
-    fs::DirEntry,
-    io::Error,
-    path::PathBuf,
-};
-
 use crate::{
     known::Environment,
-    locator::Locator,
-    messaging::{MessageDispatcher, PythonEnvironment},
+    locator::{Locator, LocatorResult},
+    messaging::PythonEnvironment,
     utils::PythonEnv,
 };
 use regex::Regex;
+use std::{collections::HashSet, fs::DirEntry, io::Error, path::PathBuf};
 
 fn is_symlinked_python_executable(path: Result<DirEntry, Error>) -> Option<PathBuf> {
     let path = path.ok()?.path();
@@ -30,37 +24,28 @@ fn is_symlinked_python_executable(path: Result<DirEntry, Error>) -> Option<PathB
 }
 
 pub struct Homebrew<'a> {
-    pub environments: HashMap<String, PythonEnvironment>,
     pub environment: &'a dyn Environment,
 }
 
 impl Homebrew<'_> {
     pub fn with<'a>(environment: &'a impl Environment) -> Homebrew {
-        Homebrew {
-            environments: HashMap::new(),
-            environment,
-        }
+        Homebrew { environment }
     }
 }
 
 impl Locator for Homebrew<'_> {
-    fn is_known(&self, python_executable: &PathBuf) -> bool {
-        self.environments
-            .contains_key(python_executable.to_str().unwrap_or_default())
+    fn resolve(&self, _env: &PythonEnv) -> Option<PythonEnvironment> {
+        None
     }
 
-    fn track_if_compatible(&mut self, _env: &PythonEnv) -> bool {
-        // We will find everything in gather
-        false
-    }
-
-    fn gather(&mut self) -> Option<()> {
+    fn find(&self) -> Option<LocatorResult> {
         let homebrew_prefix = self
             .environment
             .get_env_var("HOMEBREW_PREFIX".to_string())?;
         let homebrew_prefix_bin = PathBuf::from(homebrew_prefix).join("bin");
         let mut reported: HashSet<String> = HashSet::new();
         let python_regex = Regex::new(r"/(\d+\.\d+\.\d+)/").unwrap();
+        let mut environments: Vec<PythonEnvironment> = vec![];
         for file in std::fs::read_dir(homebrew_prefix_bin).ok()? {
             if let Some(exe) = is_symlinked_python_executable(file) {
                 let python_version = exe.to_string_lossy().to_string();
@@ -85,16 +70,13 @@ impl Locator for Homebrew<'_> {
                     None,
                     Some(vec![exe.to_string_lossy().to_string()]),
                 );
-                self.environments
-                    .insert(exe.to_string_lossy().to_string(), env);
+                environments.push(env);
             }
         }
-        Some(())
-    }
-
-    fn report(&self, reporter: &mut dyn MessageDispatcher) {
-        for env in self.environments.values() {
-            reporter.report_environment(env.clone());
+        if environments.is_empty() {
+            None
+        } else {
+            Some(LocatorResult::Environments(environments))
         }
     }
 }

native_locator/src/locator.rs

@@ -1,26 +1,27 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use crate::{messaging::MessageDispatcher, utils::PythonEnv};
-use std::path::PathBuf;
+use crate::{
+    messaging::{EnvManager, PythonEnvironment},
+    utils::PythonEnv,
+};
+
+#[derive(Debug)]
+pub enum LocatorResult {
+    Managers(Vec<EnvManager>),
+    Environments(Vec<PythonEnvironment>),
+}
 
 pub trait Locator {
     /**
-     * Whether the given Python executable is known to this locator.
-     */
-    fn is_known(&self, python_executable: &PathBuf) -> bool;
-    /**
-     * Track the given Python executable if it is compatible with the environments supported by this locator.
-     * This way, when report is called, the environment passed here will be reported as a known environment by this locator.
-     * Returns true if the environment was tracked, false otherwise.
-     */
-    fn track_if_compatible(&mut self, env: &PythonEnv) -> bool;
-    /**
-     * Finds all environments managed by this locator.
+     * Given a Python environment, this will convert it to a PythonEnvironment that can be supported by this locator.
+     * If an environment is not supported by this locator, this will return None.
+     *
+     * I.e. use this to test whether an environment is of a specific type.
      */
-    fn gather(&mut self) -> Option<()>;
+    fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment>;
     /**
-     * Report all of the tracked environments and managers.
+     * Finds all environments specific to this locator.
      */
-    fn report(&self, reporter: &mut dyn MessageDispatcher);
+    fn find(&self) -> Option<LocatorResult>;
 }