From b76efc62a72ed8031aa91ae80bb1b4a159bd3d21 Mon Sep 17 00:00:00 2001
From: markmmm <markmmm@users.noreply.github.com>
Date: Mon, 19 Feb 2024 18:09:12 +0800
Subject: [PATCH] Support dotted function paths for script entrypoints (#1622)

Co-authored-by: markm <mark.mcmahon@autodesk.com>
Co-authored-by: Micha Reiser <micha@reiser.io>
---
 crates/install-wheel-rs/src/script.rs | 24 ++++++++++++++++++++++--
 crates/install-wheel-rs/src/wheel.rs  | 16 +++++++++-------
 crates/uv/tests/pip_sync.rs           |  2 +-
 3 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/crates/install-wheel-rs/src/script.rs b/crates/install-wheel-rs/src/script.rs
index a7727dbe2..d215d0433 100644
--- a/crates/install-wheel-rs/src/script.rs
+++ b/crates/install-wheel-rs/src/script.rs
@@ -1,3 +1,4 @@
+use once_cell::sync::Lazy;
 use regex::Regex;
 use rustc_hash::FxHashSet;
 use serde::Serialize;
@@ -43,9 +44,11 @@ impl Script {
         //  between the object reference and the left square bracket, between the extra names and the square brackets and colons delimiting them,
         //  and after the right square bracket."
         // – https://packaging.python.org/en/latest/specifications/entry-points/#file-format
-        let script_regex = Regex::new(r"^(?P<module>[\w\d_\-.]+)\s*:\s*(?P<function>[\w\d_\-.]+)(?:\s*\[\s*(?P<extras>(?:[^,]+,?\s*)+)\])?\s*$").unwrap();
+        static SCRIPT_REGEX: Lazy<Regex> = Lazy::new(|| {
+            Regex::new(r"^(?P<module>[\w\d_\-.]+)\s*:\s*(?P<function>[\w\d_\-.]+)(?:\s*\[\s*(?P<extras>(?:[^,]+,?\s*)+)\])?\s*$").unwrap()
+        });
 
-        let captures = script_regex
+        let captures = SCRIPT_REGEX
             .captures(value)
             .ok_or_else(|| Error::InvalidWheel(format!("invalid console script: '{value}'")))?;
         if let Some(script_extras) = captures.name("extras") {
@@ -67,6 +70,12 @@ impl Script {
             function: captures.name("function").unwrap().as_str().to_string(),
         }))
     }
+
+    pub fn import_name(&self) -> &str {
+        self.function
+            .split_once('.')
+            .map_or(&self.function, |(import_name, _)| import_name)
+    }
 }
 
 #[cfg(test)]
@@ -98,4 +107,15 @@ mod test {
             );
         }
     }
+
+    #[test]
+    fn test_split_of_import_name_from_function() {
+        let entrypoint = "foomod:mod_bar.sub_foo.func_baz";
+
+        let script = Script::from_value("script", entrypoint, None)
+            .unwrap()
+            .unwrap();
+        assert_eq!(script.function, "mod_bar.sub_foo.func_baz");
+        assert_eq!(script.import_name(), "mod_bar");
+    }
 }
diff --git a/crates/install-wheel-rs/src/wheel.rs b/crates/install-wheel-rs/src/wheel.rs
index 6140493be..fa6ca6dcd 100644
--- a/crates/install-wheel-rs/src/wheel.rs
+++ b/crates/install-wheel-rs/src/wheel.rs
@@ -52,7 +52,13 @@ const LAUNCHER_AARCH64_CONSOLE: &[u8] =
 /// Wrapper script template function
 ///
 /// <https://github.com/pypa/pip/blob/7f8a6844037fb7255cfd0d34ff8e8cf44f2598d4/src/pip/_vendor/distlib/scripts.py#L41-L48>
-fn get_script_launcher(module: &str, import_name: &str, shebang: &str) -> String {
+fn get_script_launcher(entry_point: &Script, shebang: &str) -> String {
+    let Script {
+        module, function, ..
+    } = entry_point;
+
+    let import_name = entry_point.import_name();
+
     format!(
         r##"{shebang}
 # -*- coding: utf-8 -*-
@@ -61,7 +67,7 @@ import sys
 from {module} import {import_name}
 if __name__ == "__main__":
     sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
-    sys.exit({import_name}())
+    sys.exit({function}())
 "##
     )
 }
@@ -380,11 +386,7 @@ pub(crate) fn write_script_entrypoints(
         };
 
         // Generate the launcher script.
-        let launcher_python_script = get_script_launcher(
-            &entrypoint.module,
-            &entrypoint.function,
-            &get_shebang(location),
-        );
+        let launcher_python_script = get_script_launcher(entrypoint, &get_shebang(location));
 
         // If necessary, wrap the launcher script in a Windows launcher binary.
         if cfg!(windows) {
diff --git a/crates/uv/tests/pip_sync.rs b/crates/uv/tests/pip_sync.rs
index 1a66052ac..53d980d40 100644
--- a/crates/uv/tests/pip_sync.rs
+++ b/crates/uv/tests/pip_sync.rs
@@ -2682,7 +2682,7 @@ fn repeat_requirement() -> Result<()> {
     Resolved 1 package in [TIME]
     Downloaded 1 package in [TIME]
     Installed 1 package in [TIME]
-     + anyio==4.2.0
+     + anyio==4.3.0
     "###);
 
     Ok(())