fix Support single-file scripts with PEP 723 metadata #3176

crates/pyrefly_config/src/environment/environment.rs

@@ -131,34 +131,64 @@ impl PythonEnvironment {
             );
         }
 
-        let script = "\
+        let mut command = Command::new(interpreter);
+        command.arg("-c");
+        command.arg(Self::query_script());
+
+        Self::get_env_from_command(command, interpreter.display().to_string())
+    }
+
+    /// Query the environment that `uv` would create for a PEP 723 script.
+    pub fn get_env_from_uv_script(script: &Path) -> anyhow::Result<PythonEnvironment> {
+        let Some(script_dir) = script.parent() else {
+            return Err(anyhow!(
+                "Unable to query PEP 723 environment for `{}` because it has no parent directory",
+                script.display()
+            ));
+        };
+
+        let mut command = Command::new("uv");
+        command.arg("run");
+        command.arg("--directory");
+        command.arg(script_dir);
+        command.arg("--no-project");
+        command.arg("--isolated");
+        command.arg("--with-requirements");
+        command.arg(script);
+        command.arg("python");
+        command.arg("-c");
+        command.arg(Self::query_script());
+
+        Self::get_env_from_command(command, format!("uv PEP 723 script {}", script.display()))
+    }
+
+    fn query_script() -> &'static str {
+        "\
 import json, sys, sysconfig
 platform = sys.platform
 v = sys.version_info
 version = '{}.{}.{}'.format(v.major, v.minor, v.micro)
 stdlib_paths = [p for p in [sysconfig.get_path('stdlib')] if p is not None]
 site_package_path = [p for p in sys.path if p != '' and '.zip' not in p and not p.endswith('/lib-dynload') and p not in stdlib_paths]
 print(json.dumps({'python_platform': platform, 'python_version': version, 'site_package_path': site_package_path, 'stdlib_paths': stdlib_paths}))
-";
-
-        let mut command = Command::new(interpreter);
-        command.arg("-c");
-        command.arg(script);
+"
+    }
 
+    fn get_env_from_command(
+        mut command: Command,
+        command_description: String,
+    ) -> anyhow::Result<PythonEnvironment> {
         let python_info = command.output()?;
 
         let stdout = String::from_utf8(python_info.stdout).with_context(|| {
-            format!(
-                "while parsing Python interpreter (`{}`) stdout for environment configuration",
-                interpreter.display()
-            )
+            format!("while parsing `{command_description}` stdout for environment configuration")
         })?;
         if !python_info.status.success() {
             let stderr = String::from_utf8(python_info.stderr)
                 .unwrap_or("<Failed to parse STDOUT from UTF-8 string>".to_owned());
             return Err(anyhow::anyhow!(
-                "Unable to query interpreter {} for environment info:\nSTDOUT: {}\nSTDERR: {}",
-                interpreter.display(),
+                "Unable to query {} for environment info:\nSTDOUT: {}\nSTDERR: {}",
+                command_description,
                 stdout,
                 stderr
             ));

crates/pyrefly_config/src/environment/interpreters.rs

@@ -98,6 +98,14 @@ impl Interpreters {
         self.python_interpreter_path = Some(ConfigOrigin::lsp(interpreter));
     }
 
+    /// Freeze interpreter discovery and clear any previously selected interpreter source.
+    pub fn disable_query(&mut self) {
+        self.skip_interpreter_query = true;
+        self.python_interpreter_path = None;
+        self.fallback_python_interpreter_name = None;
+        self.conda_environment = None;
+    }
+
     /// Finds interpreters by searching in prioritized locations for the given project
     /// and interpreter settings.
     ///

crates/pyrefly_config/src/finder.rs

@@ -103,6 +103,13 @@ type LoadCallback = Box<dyn Fn(&Path) -> (ArcId<ConfigFile>, Vec<ConfigError>) +
 type FallbackCallback =
     Box<dyn Fn(ModuleNameWithKind, &ModulePath) -> ArcId<ConfigFile> + Send + Sync>;
 
+/// Function to augment a config for a specific Python file after it has been selected.
+type AfterCallback = Box<
+    dyn Fn(ModuleNameWithKind, &ModulePath, ArcId<ConfigFile>) -> anyhow::Result<ArcId<ConfigFile>>
+        + Send
+        + Sync,
+>;
+
 /// A way to find a config file given a directory or Python file.
 /// Uses a lot of caching.
 pub struct ConfigFinder {
@@ -113,6 +120,8 @@ pub struct ConfigFinder {
     /// If this returns a value, it is _not_ cached.
     /// If this returns anything other than `Ok`, the rest of the functions are used.
     before: BeforeCallback,
+    /// If this returns `Err`, the original config is returned and the error is recorded.
+    after: AfterCallback,
     /// If there is no config file, or loading it fails, use this fallback.
     fallback: FallbackCallback,
     clear_extra_caches: Box<dyn Fn() + Send + Sync>,
@@ -128,6 +137,7 @@ impl ConfigFinder {
         Self::new_custom(
             Box::new(|_, _| Ok(None)),
             load,
+            Box::new(|_, _, config| Ok(config)),
             fallback,
             clear_extra_caches,
         )
@@ -142,6 +152,7 @@ impl ConfigFinder {
         Self::new_custom(
             Box::new(move |_, _| Ok(Some(c1.dupe()))),
             Box::new(move |_| (c2.dupe(), Vec::new())),
+            Box::new(move |_, _, config| Ok(config)),
             Box::new(move |_, _| c3.dupe()),
             Box::new(|| {}),
         )
@@ -153,6 +164,7 @@ impl ConfigFinder {
     pub fn new_custom(
         before: BeforeCallback,
         load: LoadCallback,
+        after: AfterCallback,
         fallback: FallbackCallback,
         clear_extra_caches: Box<dyn Fn() + Send + Sync>,
     ) -> Self {
@@ -208,6 +220,7 @@ impl ConfigFinder {
             ),
             errors,
             before,
+            after,
             fallback,
             clear_extra_caches,
         }
@@ -264,7 +277,7 @@ impl ConfigFinder {
             None => (self.fallback)(name, path),
         };
 
-        match path.details() {
+        let config = match path.details() {
             ModulePathDetails::FileSystem(x) | ModulePathDetails::Memory(x) => {
                 let absolute = x.absolutize();
                 f(absolute.parent())
@@ -273,6 +286,14 @@ impl ConfigFinder {
             ModulePathDetails::BundledTypeshed(_) => f(None),
             ModulePathDetails::BundledTypeshedThirdParty(_) => f(None),
             ModulePathDetails::BundledThirdParty(_) => f(None),
+        };
+
+        match (self.after)(name, path, config.dupe()) {
+            Ok(config) => config,
+            Err(e) => {
+                self.errors.lock().push(ConfigError::warn(e));
+                config
+            }
         }
     }