fix: formalize refresh locator state (Fixes #387)

Author: karthiknadig

crates/pet-conda/src/lib.rs

@@ -16,7 +16,7 @@ use pet_core::{
     os_environment::Environment,
     python_environment::{PythonEnvironment, PythonEnvironmentKind},
     reporter::Reporter,
-    Locator, LocatorKind,
+    Locator, LocatorKind, RefreshStatePersistence, RefreshStateSyncScope,
 };
 use pet_fs::path::norm_case;
 use rayon::prelude::*;
@@ -216,11 +216,39 @@ impl Locator for Conda {
     fn get_kind(&self) -> LocatorKind {
         LocatorKind::Conda
     }
-    fn configure(&self, config: &pet_core::Configuration) {
-        if let Some(ref conda_exe) = config.conda_executable {
-            let mut conda_executable = self.conda_executable.write().unwrap();
-            conda_executable.replace(conda_exe.clone());
+    fn refresh_state(&self) -> RefreshStatePersistence {
+        RefreshStatePersistence::SyncedDiscoveryState
+    }
+    fn sync_refresh_state_from(&self, source: &dyn Locator, scope: &RefreshStateSyncScope) {
+        let source = source.as_any().downcast_ref::<Conda>().unwrap_or_else(|| {
+            panic!("attempted to sync Conda state from {:?}", source.get_kind())
+        });
+
+        match scope {
+            RefreshStateSyncScope::Full => {}
+            RefreshStateSyncScope::GlobalFiltered(kind)
+                if self.supported_categories().contains(kind) => {}
+            RefreshStateSyncScope::GlobalFiltered(_) | RefreshStateSyncScope::Workspace => {
+                return;
+            }
         }
+
+        self.environments.clear();
+        self.environments
+            .insert_many(source.environments.clone_map());
+
+        self.managers.clear();
+        self.managers.insert_many(source.managers.clone_map());
+
+        self.mamba_managers.clear();
+        self.mamba_managers
+            .insert_many(source.mamba_managers.clone_map());
+    }
+    fn configure(&self, config: &pet_core::Configuration) {
+        self.conda_executable
+            .write()
+            .unwrap()
+            .clone_from(&config.conda_executable);
     }
     fn supported_categories(&self) -> Vec<PythonEnvironmentKind> {
         vec![PythonEnvironmentKind::Conda]

crates/pet-core/src/lib.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use std::path::PathBuf;
+use std::{any::Any, path::PathBuf};
 
 use env::PythonEnv;
 use manager::EnvManager;
@@ -61,7 +61,26 @@ pub enum LocatorKind {
     WindowsStore,
 }
 
-pub trait Locator: Send + Sync {
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum RefreshStatePersistence {
+    /// The locator keeps no mutable state across requests.
+    Stateless,
+    /// The locator keeps configured inputs only; refresh must not copy them back.
+    ConfiguredOnly,
+    /// The locator keeps cache-like state, but later requests can repopulate it on demand.
+    SelfHydratingCache,
+    /// The locator keeps refresh-discovered state that later requests depend on for correctness.
+    SyncedDiscoveryState,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RefreshStateSyncScope {
+    Full,
+    GlobalFiltered(PythonEnvironmentKind),
+    Workspace,
+}
+
+pub trait Locator: Any + Send + Sync {
     /// Returns the name of the locator.
     fn get_kind(&self) -> LocatorKind;
     /// Configures the locator with the given configuration.
@@ -100,6 +119,21 @@ pub trait Locator: Send + Sync {
     fn configure(&self, _config: &Configuration) {
         //
     }
+    /// Describes what mutable state, if any, must survive a refresh boundary.
+    ///
+    /// Refresh runs execute against transient locator graphs and then invoke
+    /// `sync_refresh_state_from()` on the long-lived shared locators.
+    fn refresh_state(&self) -> RefreshStatePersistence {
+        RefreshStatePersistence::Stateless
+    }
+    /// Copies correctness-critical post-refresh state from a transient locator into this
+    /// long-lived shared locator.
+    ///
+    /// Override this only when `refresh_state()` returns
+    /// `RefreshStatePersistence::SyncedDiscoveryState`.
+    fn sync_refresh_state_from(&self, _source: &dyn Locator, _scope: &RefreshStateSyncScope) {
+        //
+    }
     /// Returns a list of supported categories for this locator.
     fn supported_categories(&self) -> Vec<PythonEnvironmentKind>;
     /// Given a Python executable, and some optional data like prefix,
@@ -112,3 +146,9 @@ pub trait Locator: Send + Sync {
     /// Finds all environments specific to this locator.
     fn find(&self, reporter: &dyn Reporter);
 }
+
+impl dyn Locator {
+    pub fn as_any(&self) -> &dyn Any {
+        self
+    }
+}

crates/pet-linux-global-python/src/lib.rs

@@ -15,7 +15,7 @@ use pet_core::{
     env::PythonEnv,
     python_environment::{PythonEnvironment, PythonEnvironmentBuilder, PythonEnvironmentKind},
     reporter::Reporter,
-    Locator, LocatorKind,
+    Locator, LocatorKind, RefreshStatePersistence,
 };
 use pet_fs::path::resolve_symlink;
 use pet_python_utils::{env::ResolvedPythonEnv, executable::find_executables};
@@ -62,6 +62,9 @@ impl Locator for LinuxGlobalPython {
     fn get_kind(&self) -> LocatorKind {
         LocatorKind::LinuxGlobal
     }
+    fn refresh_state(&self) -> RefreshStatePersistence {
+        RefreshStatePersistence::SelfHydratingCache
+    }
     fn supported_categories(&self) -> Vec<PythonEnvironmentKind> {
         vec![PythonEnvironmentKind::LinuxGlobal]
     }

crates/pet-pipenv/src/lib.rs

@@ -11,7 +11,7 @@ use pet_core::LocatorKind;
 use pet_core::{
     python_environment::{PythonEnvironment, PythonEnvironmentBuilder, PythonEnvironmentKind},
     reporter::Reporter,
-    Configuration, Locator,
+    Configuration, Locator, RefreshStatePersistence,
 };
 use pet_fs::path::norm_case;
 use pet_python_utils::executable::find_executables;
@@ -418,10 +418,15 @@ impl Locator for PipEnv {
         LocatorKind::PipEnv
     }
 
+    fn refresh_state(&self) -> RefreshStatePersistence {
+        RefreshStatePersistence::ConfiguredOnly
+    }
+
     fn configure(&self, config: &Configuration) {
-        if let Some(exe) = &config.pipenv_executable {
-            self.pipenv_executable.write().unwrap().replace(exe.clone());
-        }
+        self.pipenv_executable
+            .write()
+            .unwrap()
+            .clone_from(&config.pipenv_executable);
     }
 
     fn supported_categories(&self) -> Vec<PythonEnvironmentKind> {

crates/pet-poetry/src/lib.rs

@@ -11,7 +11,8 @@ use pet_core::{
     os_environment::Environment,
     python_environment::{PythonEnvironment, PythonEnvironmentKind},
     reporter::Reporter,
-    Configuration, Locator, LocatorKind, LocatorResult,
+    Configuration, Locator, LocatorKind, LocatorResult, RefreshStatePersistence,
+    RefreshStateSyncScope,
 };
 use pet_virtualenv::is_virtualenv;
 use regex::Regex;
@@ -137,7 +138,6 @@ impl Poetry {
         }
     }
     fn clear(&self) {
-        self.poetry_executable.write().unwrap().take();
         self.search_result.write().unwrap().take();
     }
     pub fn from(environment: &dyn Environment) -> Poetry {
@@ -152,6 +152,34 @@ impl Poetry {
             .clone_from(&search_result);
     }
 
+    pub fn merge_search_result_from(&self, source: &Poetry) {
+        let source_search_result = source.search_result.read().unwrap().clone();
+        let Some(source_search_result) = source_search_result else {
+            return;
+        };
+
+        let mut merged = self
+            .search_result
+            .read()
+            .unwrap()
+            .clone()
+            .unwrap_or(LocatorResult {
+                managers: vec![],
+                environments: vec![],
+            });
+        merged.managers.extend(source_search_result.managers);
+        merged.managers.sort();
+        merged.managers.dedup();
+
+        merged
+            .environments
+            .extend(source_search_result.environments);
+        merged.environments.sort();
+        merged.environments.dedup();
+
+        self.search_result.write().unwrap().replace(merged);
+    }
+
     fn find_with_cache(&self) -> Option<LocatorResult> {
         // First check if we have cached results
         {
@@ -226,17 +254,39 @@ impl Locator for Poetry {
     fn get_kind(&self) -> LocatorKind {
         LocatorKind::Poetry
     }
+    fn refresh_state(&self) -> RefreshStatePersistence {
+        RefreshStatePersistence::SyncedDiscoveryState
+    }
+    fn sync_refresh_state_from(&self, source: &dyn Locator, scope: &RefreshStateSyncScope) {
+        let source = source.as_any().downcast_ref::<Poetry>().unwrap_or_else(|| {
+            panic!(
+                "attempted to sync Poetry state from {:?}",
+                source.get_kind()
+            )
+        });
+        match scope {
+            RefreshStateSyncScope::Full => self.sync_search_result_from(source),
+            RefreshStateSyncScope::GlobalFiltered(kind)
+                if self.supported_categories().contains(kind) =>
+            {
+                self.sync_search_result_from(source)
+            }
+            RefreshStateSyncScope::Workspace => self.merge_search_result_from(source),
+            RefreshStateSyncScope::GlobalFiltered(_) => {}
+        }
+    }
     fn configure(&self, config: &Configuration) {
+        let mut ws_dirs = self.workspace_directories.write().unwrap();
+        ws_dirs.clear();
         if let Some(workspace_directories) = &config.workspace_directories {
-            let mut ws_dirs = self.workspace_directories.write().unwrap();
-            ws_dirs.clear();
             if !workspace_directories.is_empty() {
                 ws_dirs.extend(workspace_directories.clone());
             }
         }
-        if let Some(exe) = &config.poetry_executable {
-            self.poetry_executable.write().unwrap().replace(exe.clone());
-        }
+        self.poetry_executable
+            .write()
+            .unwrap()
+            .clone_from(&config.poetry_executable);
     }
 
     fn supported_categories(&self) -> Vec<PythonEnvironmentKind> {
@@ -349,4 +399,91 @@ mod tests {
             Some("fresh")
         );
     }
+
+    #[test]
+    fn test_workspace_scope_merges_search_results() {
+        let environment = EnvironmentApi::new();
+        let target = Poetry::from(&environment);
+        let source = Poetry::from(&environment);
+
+        target
+            .search_result
+            .write()
+            .unwrap()
+            .replace(LocatorResult {
+                managers: vec![],
+                environments: vec![PythonEnvironment {
+                    name: Some("existing".to_string()),
+                    kind: Some(PythonEnvironmentKind::Poetry),
+                    ..Default::default()
+                }],
+            });
+
+        source
+            .search_result
+            .write()
+            .unwrap()
+            .replace(LocatorResult {
+                managers: vec![],
+                environments: vec![PythonEnvironment {
+                    name: Some("workspace".to_string()),
+                    kind: Some(PythonEnvironmentKind::Poetry),
+                    ..Default::default()
+                }],
+            });
+
+        target.sync_refresh_state_from(&source, &RefreshStateSyncScope::Workspace);
+
+        let result = target.search_result.read().unwrap().clone().unwrap();
+        let mut names = result
+            .environments
+            .iter()
+            .map(|environment| environment.name.clone().unwrap())
+            .collect::<Vec<String>>();
+        names.sort();
+
+        assert_eq!(names, vec!["existing".to_string(), "workspace".to_string()]);
+    }
+
+    #[test]
+    fn test_clear_preserves_configured_poetry_executable() {
+        let environment = EnvironmentApi::new();
+        let poetry = Poetry::from(&environment);
+        let configured = PathBuf::from("/configured/poetry");
+
+        poetry.configure(&Configuration {
+            poetry_executable: Some(configured.clone()),
+            ..Default::default()
+        });
+        poetry
+            .search_result
+            .write()
+            .unwrap()
+            .replace(LocatorResult {
+                managers: vec![],
+                environments: vec![],
+            });
+
+        poetry.clear();
+
+        assert_eq!(
+            poetry.poetry_executable.read().unwrap().clone(),
+            Some(configured)
+        );
+        assert!(poetry.search_result.read().unwrap().is_none());
+    }
+
+    #[test]
+    fn test_configure_clears_poetry_executable_when_unset() {
+        let environment = EnvironmentApi::new();
+        let poetry = Poetry::from(&environment);
+
+        poetry.configure(&Configuration {
+            poetry_executable: Some(PathBuf::from("/configured/poetry")),
+            ..Default::default()
+        });
+        poetry.configure(&Configuration::default());
+
+        assert!(poetry.poetry_executable.read().unwrap().is_none());
+    }
 }

crates/pet-pyenv/src/lib.rs

@@ -19,7 +19,7 @@ use pet_core::{
     os_environment::Environment,
     python_environment::{PythonEnvironment, PythonEnvironmentKind},
     reporter::Reporter,
-    Locator, LocatorKind,
+    Locator, LocatorKind, RefreshStatePersistence,
 };
 use pet_python_utils::executable::find_executable;
 
@@ -84,6 +84,9 @@ impl Locator for PyEnv {
     fn get_kind(&self) -> LocatorKind {
         LocatorKind::PyEnv
     }
+    fn refresh_state(&self) -> RefreshStatePersistence {
+        RefreshStatePersistence::SelfHydratingCache
+    }
     fn supported_categories(&self) -> Vec<PythonEnvironmentKind> {
         vec![
             PythonEnvironmentKind::Pyenv,