Teach ty check to ask uv to sync the venv of a PEP-723 script

Author: Gankra

crates/ruff_python_ast/src/script.rs

@@ -12,11 +12,11 @@ static FINDER: LazyLock<Finder> = LazyLock::new(|| Finder::new(b"# /// script"))
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct ScriptTag {
     /// The content of the script before the metadata block.
-    prelude: String,
+    pub prelude: String,
     /// The metadata block.
-    metadata: String,
+    pub metadata: String,
     /// The content of the script after the metadata block.
-    postlude: String,
+    pub postlude: String,
 }
 
 impl ScriptTag {

crates/ty/src/lib.rs

@@ -90,6 +90,22 @@ fn run_check(args: CheckCommand) -> anyhow::Result<ExitStatus> {
             })?
     };
 
+    let system = OsSystem::new(&cwd);
+
+    // If we see a single path, check if it's a PEP-723 script
+    let mut script_project = None;
+    if let [path] = &*args.paths {
+        match ProjectMetadata::discover_script(path, &system) {
+            Ok(project) => {
+                script_project = Some(project);
+            }
+            Err(ty_project::ProjectMetadataError::NotAScript(_)) => {
+                // This is fine
+            }
+            Err(e) => tracing::info!("Issue reading script at `{path}`: {e}"),
+        }
+    }
+
     let project_path = args
         .project
         .as_ref()
@@ -111,7 +127,6 @@ fn run_check(args: CheckCommand) -> anyhow::Result<ExitStatus> {
         .map(|path| SystemPath::absolute(path, &cwd))
         .collect();
 
-    let system = OsSystem::new(&cwd);
     let watch = args.watch;
     let exit_zero = args.exit_zero;
     let config_file = args
@@ -121,7 +136,13 @@ fn run_check(args: CheckCommand) -> anyhow::Result<ExitStatus> {
 
     let mut project_metadata = match &config_file {
         Some(config_file) => ProjectMetadata::from_config_file(config_file.clone(), &system)?,
-        None => ProjectMetadata::discover(&project_path, &system)?,
+        None => {
+            if let Some(project) = script_project {
+                project
+            } else {
+                ProjectMetadata::discover(&project_path, &system)?
+            }
+        }
     };
 
     project_metadata.apply_configuration_files(&system)?;

crates/ty_project/src/metadata.rs

@@ -9,13 +9,15 @@ use ty_python_semantic::ProgramSettings;
 
 use crate::metadata::options::ProjectOptionsOverrides;
 use crate::metadata::pyproject::{Project, PyProject, PyProjectError, ResolveRequiresPythonError};
-use crate::metadata::value::ValueSource;
+use crate::metadata::script::{Pep723Error, Pep723Metadata};
+use crate::metadata::value::{RelativePathBuf, ValueSource};
 pub use options::Options;
 use options::TyTomlError;
 
 mod configuration_file;
 pub mod options;
 pub mod pyproject;
+pub mod script;
 pub mod settings;
 pub mod value;
 
@@ -85,6 +87,32 @@ impl ProjectMetadata {
         )
     }
 
+    /// Loads a project from a `pyproject.toml` file.
+    pub(crate) fn from_script(
+        script: Pep723Metadata,
+        script_path: &SystemPath,
+    ) -> Result<Self, ResolveRequiresPythonError> {
+        let project = Some(&script.to_project());
+        let parent_dir = script_path
+            .parent()
+            .map(ToOwned::to_owned)
+            .unwrap_or_default();
+        let mut metadata = Self::from_options(
+            script.tool.and_then(|tool| tool.ty).unwrap_or_default(),
+            parent_dir,
+            project,
+        )?;
+
+        // Try to get `uv sync --script` to setup the venv for us
+        if let Some(python) = script::uv_sync_script(script_path) {
+            let mut environment = metadata.options.environment.unwrap_or_default();
+            environment.python = Some(RelativePathBuf::new(python, ValueSource::Cli));
+            metadata.options.environment = Some(environment);
+        }
+
+        Ok(metadata)
+    }
+
     /// Loads a project from a set of options with an optional pyproject-project table.
     pub fn from_options(
         mut options: Options,
@@ -120,6 +148,46 @@ impl ProjectMetadata {
         })
     }
 
+    pub fn discover_script(
+        path: &SystemPath,
+        system: &dyn System,
+    ) -> Result<ProjectMetadata, ProjectMetadataError> {
+        tracing::debug!("Searching for a PEP-723 Script in '{path}'");
+        if !system.is_file(path) {
+            return Err(ProjectMetadataError::NotAScript(path.to_path_buf()));
+        }
+
+        let script_metadata = if let Ok(script_str) = system.read_to_string(path) {
+            match Pep723Metadata::from_script_str(
+                script_str.as_bytes(),
+                ValueSource::File(Arc::new(path.to_owned())),
+            ) {
+                Ok(Some(pyproject)) => Some(pyproject),
+                Ok(None) => None,
+                Err(error) => {
+                    return Err(ProjectMetadataError::InvalidScript {
+                        path: path.to_owned(),
+                        source: Box::new(error),
+                    });
+                }
+            }
+        } else {
+            None
+        };
+
+        let Some(script_metadata) = script_metadata else {
+            return Err(ProjectMetadataError::NotAScript(path.to_path_buf()));
+        };
+
+        let metadata = ProjectMetadata::from_script(script_metadata, path).map_err(|err| {
+            ProjectMetadataError::InvalidRequiresPythonConstraint {
+                source: err,
+                path: path.to_owned(),
+            }
+        })?;
+
+        Ok(metadata)
+    }
     /// Discovers the closest project at `path` and returns its metadata.
     ///
     /// The algorithm traverses upwards in the `path`'s ancestor chain and uses the following precedence
@@ -319,12 +387,21 @@ pub enum ProjectMetadataError {
     #[error("project path '{0}' is not a directory")]
     NotADirectory(SystemPathBuf),
 
+    #[error("project path '{0}' is not a PEP-723 script")]
+    NotAScript(SystemPathBuf),
+
     #[error("{path} is not a valid `pyproject.toml`: {source}")]
     InvalidPyProject {
         source: Box<PyProjectError>,
         path: SystemPathBuf,
     },
 
+    #[error("{path} is not a valid PEP-723 script: {source}")]
+    InvalidScript {
+        source: Box<Pep723Error>,
+        path: SystemPathBuf,
+    },
+
     #[error("{path} is not a valid `ty.toml`: {source}")]
     InvalidTyToml {
         source: Box<TyTomlError>,

crates/ty_project/src/metadata/script.rs

@@ -0,0 +1,157 @@
+use std::{io, process::Command, str::FromStr};
+
+use camino::Utf8PathBuf;
+use pep440_rs::VersionSpecifiers;
+use ruff_db::system::{SystemPath, SystemPathBuf};
+use ruff_python_ast::script::ScriptTag;
+use serde::Deserialize;
+use thiserror::Error;
+
+use crate::metadata::{
+    pyproject::{Project, Tool},
+    value::{RangedValue, ValueSource, ValueSourceGuard},
+};
+
+/// PEP 723 metadata as parsed from a `script` comment block.
+///
+/// See: <https://peps.python.org/pep-0723/>
+#[derive(Debug, Deserialize, Clone)]
+#[serde(rename_all = "kebab-case")]
+pub struct Pep723Metadata {
+    pub dependencies: Option<RangedValue<Vec<toml::Value>>>,
+    pub requires_python: Option<RangedValue<VersionSpecifiers>>,
+    pub tool: Option<Tool>,
+
+    /// The raw unserialized document.
+    #[serde(skip)]
+    pub raw: String,
+}
+
+#[derive(Debug, Error)]
+pub enum Pep723Error {
+    #[error(
+        "An opening tag (`# /// script`) was found without a closing tag (`# ///`). Ensure that every line between the opening and closing tags (including empty lines) starts with a leading `#`."
+    )]
+    UnclosedBlock,
+    #[error("The PEP 723 metadata block is missing from the script.")]
+    MissingTag,
+    #[error(transparent)]
+    Io(#[from] io::Error),
+    #[error(transparent)]
+    Utf8(#[from] std::str::Utf8Error),
+    #[error(transparent)]
+    Toml(#[from] toml::de::Error),
+    #[error("Invalid filename `{0}` supplied")]
+    InvalidFilename(String),
+}
+
+impl Pep723Metadata {
+    /// Parse the PEP 723 metadata from `stdin`.
+    pub fn from_script_str(
+        contents: &[u8],
+        source: ValueSource,
+    ) -> Result<Option<Self>, Pep723Error> {
+        let _guard = ValueSourceGuard::new(source, true);
+
+        // Extract the `script` tag.
+        let Some(ScriptTag { metadata, .. }) = ScriptTag::parse(contents) else {
+            return Ok(None);
+        };
+
+        // Parse the metadata.
+        Ok(Some(Self::from_str(&metadata)?))
+    }
+
+    pub fn to_project(&self) -> Project {
+        Project {
+            name: None,
+            version: None,
+            requires_python: self.requires_python.clone(),
+        }
+    }
+}
+
+/*
+{
+  "schema": {
+    "version": "preview"
+  },
+  "target": "script",
+  "script": {
+    "path": "/Users/myuser/code/myproj/scripts/load-test.py"
+  },
+  "sync": {
+    "environment": {
+      "path": "/Users/myuser/.cache/uv/environments-v2/load-test-d6edaf5bfab110a8",
+      "python": {
+        "path": "/Users/myuser/.cache/uv/environments-v2/load-test-d6edaf5bfab110a8/bin/python3",
+        "version": "3.14.0",
+        "implementation": "cpython"
+      }
+    },
+    "action": "check"
+  },
+  "lock": null,
+  "dry_run": false
+}
+*/
+
+/// The output of `uv sync --output-format=json --script ...`
+#[derive(Debug, Clone, Deserialize)]
+struct UvMetadata {
+    sync: Option<UvSync>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+struct UvSync {
+    environment: Option<UvEnvironment>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+struct UvEnvironment {
+    path: Option<String>,
+}
+
+/// Ask `uv` to sync the script's venv to some temp dir so we can analyze dependencies properly
+///
+/// Returns the path to the venv on success
+pub fn uv_sync_script(script_path: &SystemPath) -> Option<SystemPathBuf> {
+    tracing::info!("Asking uv to sync the script's venv");
+    let mut command = Command::new("uv");
+    command
+        .arg("sync")
+        .arg("--output-format=json")
+        .arg("--script")
+        .arg(script_path.as_str());
+    let output = command
+        .output()
+        .inspect_err(|e| {
+            tracing::info!(
+                "failed to run `uv sync --output-format=json --script {script_path}`: {e}"
+            );
+        })
+        .ok()?;
+    let metadata: UvMetadata = serde_json::from_slice(&output.stdout)
+        .inspect_err(|e| {
+            tracing::info!(
+                "failed to parse `uv sync --output-format=json --script {script_path}`: {e}"
+            );
+        })
+        .ok()?;
+    let env_path = metadata.sync?.environment?.path?;
+    let utf8_path = Utf8PathBuf::from(env_path);
+    Some(SystemPathBuf::from_utf8_path_buf(utf8_path))
+}
+
+impl FromStr for Pep723Metadata {
+    type Err = toml::de::Error;
+
+    /// Parse `Pep723Metadata` from a raw TOML string.
+    fn from_str(raw: &str) -> Result<Self, Self::Err> {
+        let metadata = toml::from_str(raw)?;
+        Ok(Self {
+            raw: raw.to_string(),
+            ..metadata
+        })
+    }
+}