From bb0fb8e9bf3b1ffbb9ad7ff5836038ced23e50e0 Mon Sep 17 00:00:00 2001
From: Jo <10510431+j178@users.noreply.github.com>
Date: Thu, 12 Sep 2024 03:34:29 +0800
Subject: [PATCH] `uv run` supports python zipapp (#7289)

## Summary

`python` supports running a zipfile containing a `__main__.py` file, for
example `python ./pre-commit-3.8.0.pyz`.

See https://docs.python.org/3/using/cmdline.html#interface-options:

> <script> Execute the Python code contained in script, which must be a
filesystem path (absolute or relative) referring to either a Python
file, a directory containing a __main__.py file, or a zipfile containing
a __main__.py file.

and https://docs.python.org/3/library/zipapp.html.

Similar to #7281, this PR allows `uv run ./pre-commit-3.8.0.pyz` to
work.

## Test Plan

```console
$ curl -O https://github.com/pre-commit/pre-commit/releases/download/v3.8.0/pre-commit-3.8.0.pyz
$ cargo run -- run ./pre-commit-3.8.0.pyz
```
---
 crates/uv/Cargo.toml                  |  1 +
 crates/uv/src/commands/project/run.rs | 44 +++++++++++++++++++++------
 crates/uv/tests/run.rs                | 39 ++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 9 deletions(-)

diff --git a/crates/uv/Cargo.toml b/crates/uv/Cargo.toml
index 2b65aed76..09f4c737c 100644
--- a/crates/uv/Cargo.toml
+++ b/crates/uv/Cargo.toml
@@ -79,6 +79,7 @@ tracing-tree = { workspace = true }
 unicode-width = { workspace = true }
 url = { workspace = true }
 which = { workspace = true }
+zip = { workspace = true }
 
 [target.'cfg(target_os = "windows")'.dependencies]
 mimalloc = { version = "0.1.39" }
diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs
index 064953735..8334db395 100644
--- a/crates/uv/src/commands/project/run.rs
+++ b/crates/uv/src/commands/project/run.rs
@@ -780,6 +780,9 @@ pub(crate) enum RunCommand {
     PythonGuiScript(PathBuf, Vec<OsString>),
     /// Execute a Python package containing a `__main__.py` file.
     PythonPackage(PathBuf, Vec<OsString>),
+    /// Execute a Python [zipapp].
+    /// [zipapp]: <https://docs.python.org/3/library/zipapp.html>
+    PythonZipapp(PathBuf, Vec<OsString>),
     /// Execute a `python` script provided via `stdin`.
     PythonStdin(Vec<u8>),
     /// Execute an external command.
@@ -793,10 +796,11 @@ impl RunCommand {
     fn display_executable(&self) -> Cow<'_, str> {
         match self {
             Self::Python(_) => Cow::Borrowed("python"),
-            Self::PythonScript(_, _) | Self::PythonPackage(_, _) | Self::Empty => {
-                Cow::Borrowed("python")
-            }
-            Self::PythonGuiScript(_, _) => Cow::Borrowed("pythonw"),
+            Self::PythonScript(..)
+            | Self::PythonPackage(..)
+            | Self::PythonZipapp(..)
+            | Self::Empty => Cow::Borrowed("python"),
+            Self::PythonGuiScript(..) => Cow::Borrowed("pythonw"),
             Self::PythonStdin(_) => Cow::Borrowed("python -c"),
             Self::External(executable, _) => executable.to_string_lossy(),
         }
@@ -810,7 +814,9 @@ impl RunCommand {
                 process.args(args);
                 process
             }
-            Self::PythonScript(target, args) | Self::PythonPackage(target, args) => {
+            Self::PythonScript(target, args)
+            | Self::PythonPackage(target, args)
+            | Self::PythonZipapp(target, args) => {
                 let mut process = Command::new(interpreter.sys_executable());
                 process.arg(target);
                 process.args(args);
@@ -873,7 +879,9 @@ impl std::fmt::Display for RunCommand {
                 }
                 Ok(())
             }
-            Self::PythonScript(target, args) | Self::PythonPackage(target, args) => {
+            Self::PythonScript(target, args)
+            | Self::PythonPackage(target, args)
+            | Self::PythonZipapp(target, args) => {
                 write!(f, "python {}", target.display())?;
                 for arg in args {
                     write!(f, " {}", arg.to_string_lossy())?;
@@ -917,6 +925,10 @@ impl TryFrom<&ExternalCommand> for RunCommand {
         };
 
         let target_path = PathBuf::from(&target);
+        let metadata = target_path.metadata();
+        let is_file = metadata.as_ref().map_or(false, std::fs::Metadata::is_file);
+        let is_dir = metadata.as_ref().map_or(false, std::fs::Metadata::is_dir);
+
         if target.eq_ignore_ascii_case("-") {
             let mut buf = Vec::with_capacity(1024);
             std::io::stdin().read_to_end(&mut buf)?;
@@ -926,18 +938,20 @@ impl TryFrom<&ExternalCommand> for RunCommand {
         } else if target_path
             .extension()
             .is_some_and(|ext| ext.eq_ignore_ascii_case("py") || ext.eq_ignore_ascii_case("pyc"))
-            && target_path.exists()
+            && is_file
         {
             Ok(Self::PythonScript(target_path, args.to_vec()))
         } else if cfg!(windows)
             && target_path
                 .extension()
                 .is_some_and(|ext| ext.eq_ignore_ascii_case("pyw"))
-            && target_path.exists()
+            && is_file
         {
             Ok(Self::PythonGuiScript(target_path, args.to_vec()))
-        } else if target_path.is_dir() && target_path.join("__main__.py").exists() {
+        } else if is_dir && target_path.join("__main__.py").is_file() {
             Ok(Self::PythonPackage(target_path, args.to_vec()))
+        } else if is_file && is_python_zipapp(&target_path) {
+            Ok(Self::PythonZipapp(target_path, args.to_vec()))
         } else {
             Ok(Self::External(
                 target.clone(),
@@ -946,3 +960,15 @@ impl TryFrom<&ExternalCommand> for RunCommand {
         }
     }
 }
+
+/// Returns `true` if the target is a ZIP archive containing a `__main__.py` file.
+fn is_python_zipapp(target: &Path) -> bool {
+    if let Ok(file) = fs_err::File::open(target) {
+        if let Ok(mut archive) = zip::ZipArchive::new(file) {
+            return archive
+                .by_name("__main__.py")
+                .map_or(false, |f| f.is_file());
+        }
+    }
+    false
+}
diff --git a/crates/uv/tests/run.rs b/crates/uv/tests/run.rs
index 8a35c0f89..300b41309 100644
--- a/crates/uv/tests/run.rs
+++ b/crates/uv/tests/run.rs
@@ -1656,6 +1656,45 @@ fn run_package() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn run_zipapp() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    // Create a zipapp.
+    let child = context.temp_dir.child("app");
+    child.create_dir_all()?;
+
+    let main_script = child.child("__main__.py");
+    main_script.write_str(indoc! { r#"
+        print("Hello, world!")
+       "#
+    })?;
+
+    let zipapp = context.temp_dir.child("app.pyz");
+    let status = context
+        .run()
+        .arg("python")
+        .arg("-m")
+        .arg("zipapp")
+        .arg(child.as_ref())
+        .arg("--output")
+        .arg(zipapp.as_ref())
+        .status()?;
+    assert!(status.success());
+
+    // Run the zipapp.
+    uv_snapshot!(context.filters(), context.run().arg(zipapp.as_ref()), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    Hello, world!
+
+    ----- stderr -----
+    "###);
+
+    Ok(())
+}
+
 /// When the `pyproject.toml` file is invalid.
 #[test]
 fn run_project_toml_error() -> Result<()> {