[red-knot] Include vendored typeshed stubs as a zipfile in the Ruff binary (#11779)

Co-authored-by: Micha Reiser <micha@reiser.io>
Co-authored-by: Carl Meyer <carl@astral.sh>

.github/renovate.json5

@@ -41,6 +41,13 @@
       description: "Disable PRs updating GitHub runners (e.g. 'runs-on: macos-14')",
       enabled: false,
     },
+    {
+      // Disable updates of `zip-rs`; intentionally pinned for now due to ownership change
+      // See: https://github.com/astral-sh/uv/issues/3642
+      matchPackagePatterns: ["zip"],
+      matchManagers: ["cargo"],
+      enabled: false,
+    },
     {
       groupName: "pre-commit dependencies",
       matchManagers: ["pre-commit"],

Cargo.lock

@@ -199,6 +199,12 @@ version = "3.16.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c"
 
+[[package]]
+name = "byteorder"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
+
 [[package]]
 name = "cachedir"
 version = "0.3.1"
@@ -219,6 +225,11 @@ name = "cc"
 version = "1.0.95"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d32a725bc159af97c3e629873bb9f88fb8cf8a4867175f76dc987815ea07c83b"
+dependencies = [
+ "jobserver",
+ "libc",
+ "once_cell",
+]
 
 [[package]]
 name = "cfg-if"
@@ -1123,6 +1134,15 @@ version = "1.0.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b"
 
+[[package]]
+name = "jobserver"
+version = "0.1.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2b099aaa34a9751c5bf0878add70444e1ed2dd73f347be99003d4577277de6e"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "jod-thread"
 version = "0.1.2"
@@ -1629,6 +1649,12 @@ version = "0.2.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02"
 
+[[package]]
+name = "pkg-config"
+version = "0.3.30"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec"
+
 [[package]]
 name = "portable-atomic"
 version = "1.6.0"
@@ -1782,6 +1808,8 @@ dependencies = [
  "tracing",
  "tracing-subscriber",
  "tracing-tree",
+ "walkdir",
+ "zip",
 ]
 
 [[package]]
@@ -3593,3 +3621,44 @@ name = "zeroize"
 version = "1.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "525b4ec142c6b68a2d10f01f7bbf6755599ca3f81ea53b8431b7dd348f5fdb2d"
+
+[[package]]
+name = "zip"
+version = "0.6.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261"
+dependencies = [
+ "byteorder",
+ "crc32fast",
+ "crossbeam-utils",
+ "zstd",
+]
+
+[[package]]
+name = "zstd"
+version = "0.11.2+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4"
+dependencies = [
+ "zstd-safe",
+]
+
+[[package]]
+name = "zstd-safe"
+version = "5.0.2+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d2a5585e04f9eea4b2a3d1eca508c4dee9592a89ef6f450c11719da0726f4db"
+dependencies = [
+ "libc",
+ "zstd-sys",
+]
+
+[[package]]
+name = "zstd-sys"
+version = "2.0.10+zstd.1.5.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c253a4914af5bafc8fa8c86ee400827e83cf6ec01195ec1f1ed8441bf00d65aa"
+dependencies = [
+ "cc",
+ "pkg-config",
+]

Cargo.toml

@@ -136,6 +136,7 @@ walkdir = { version = "2.3.2" }
 wasm-bindgen = { version = "0.2.92" }
 wasm-bindgen-test = { version = "0.3.42" }
 wild = { version = "2" }
+zip = { version = "0.6.6", default-features = false, features = ["zstd"] }
 
 [workspace.lints.rust]
 unsafe_code = "warn"

crates/red_knot/Cargo.toml

@@ -34,6 +34,11 @@ smol_str = { version = "0.2.1" }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 tracing-tree = { workspace = true }
+zip = { workspace = true }
+
+[build-dependencies]
+zip = { workspace = true }
+walkdir = { workspace = true }
 
 [dev-dependencies]
 tempfile = { workspace = true }

crates/red_knot/build.rs

@@ -0,0 +1,73 @@
+//! Build script to package our vendored typeshed files
+//! into a zip archive that can be included in the Ruff binary.
+//!
+//! This script should be automatically run at build time
+//! whenever the script itself changes, or whenever any files
+//! in `crates/red_knot/vendor/typeshed` change.
+
+use std::fs::File;
+use std::path::Path;
+
+use zip::result::ZipResult;
+use zip::write::{FileOptions, ZipWriter};
+use zip::CompressionMethod;
+
+const TYPESHED_SOURCE_DIR: &str = "vendor/typeshed";
+const TYPESHED_ZIP_LOCATION: &str = "/zipped_typeshed.zip";
+
+/// Recursively zip the contents of an entire directory.
+///
+/// This routine is adapted from a recipe at
+/// <https://github.com/zip-rs/zip-old/blob/5d0f198124946b7be4e5969719a7f29f363118cd/examples/write_dir.rs>
+fn zip_dir(directory_path: &str, writer: File) -> ZipResult<File> {
+    let mut zip = ZipWriter::new(writer);
+
+    let options = FileOptions::default()
+        .compression_method(CompressionMethod::Zstd)
+        .unix_permissions(0o644);
+
+    for entry in walkdir::WalkDir::new(directory_path) {
+        let dir_entry = entry.unwrap();
+        let relative_path = dir_entry.path();
+        let name = relative_path
+            .strip_prefix(Path::new(directory_path))
+            .unwrap()
+            .to_str()
+            .expect("Unexpected non-utf8 typeshed path!");
+
+        // Write file or directory explicitly
+        // Some unzip tools unzip files with directory paths correctly, some do not!
+        if relative_path.is_file() {
+            println!("adding file {relative_path:?} as {name:?} ...");
+            zip.start_file(name, options)?;
+            let mut f = File::open(relative_path)?;
+            std::io::copy(&mut f, &mut zip).unwrap();
+        } else if !name.is_empty() {
+            // Only if not root! Avoids path spec / warning
+            // and mapname conversion failed error on unzip
+            println!("adding dir {relative_path:?} as {name:?} ...");
+            zip.add_directory(name, options)?;
+        }
+    }
+    zip.finish()
+}
+
+fn main() {
+    println!("cargo:rerun-if-changed={TYPESHED_SOURCE_DIR}");
+    assert!(
+        Path::new(TYPESHED_SOURCE_DIR).is_dir(),
+        "Where is typeshed?"
+    );
+    let out_dir = std::env::var("OUT_DIR").unwrap();
+
+    // N.B. Deliberately using `format!()` instead of `Path::join()` here,
+    // so that we use `/` as a path separator on all platforms.
+    // That enables us to load the typeshed zip at compile time in `module.rs`
+    // (otherwise we'd have to dynamically determine the exact path to the typeshed zip
+    // based on the default path separator for the specific platform we're on,
+    // which can't be done at compile time.)
+    let zipped_typeshed_location = format!("{out_dir}{TYPESHED_ZIP_LOCATION}");
+
+    let zipped_typeshed = File::create(zipped_typeshed_location).unwrap();
+    zip_dir(TYPESHED_SOURCE_DIR, zipped_typeshed).unwrap();
+}

crates/red_knot/src/module.rs

@@ -770,8 +770,11 @@ impl PackageKind {
 
 #[cfg(test)]
 mod tests {
+    use std::io::{Cursor, Read};
     use std::num::NonZeroU32;
-    use std::path::PathBuf;
+    use std::path::{Path, PathBuf};
+
+    use zip::ZipArchive;
 
     use crate::db::tests::TestDb;
     use crate::db::SourceDb;
@@ -923,6 +926,28 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn typeshed_zip_created_at_build_time() -> anyhow::Result<()> {
+        // The file path here is hardcoded in this crate's `build.rs` script.
+        // Luckily this crate will fail to build if this file isn't available at build time.
+        const TYPESHED_ZIP_BYTES: &[u8] =
+            include_bytes!(concat!(env!("OUT_DIR"), "/zipped_typeshed.zip"));
+        assert!(!TYPESHED_ZIP_BYTES.is_empty());
+        let mut typeshed_zip_archive = ZipArchive::new(Cursor::new(TYPESHED_ZIP_BYTES))?;
+
+        let path_to_functools = Path::new("stdlib").join("functools.pyi");
+        let mut functools_module_stub = typeshed_zip_archive
+            .by_name(path_to_functools.to_str().unwrap())
+            .unwrap();
+        assert!(functools_module_stub.is_file());
+
+        let mut functools_module_stub_source = String::new();
+        functools_module_stub.read_to_string(&mut functools_module_stub_source)?;
+
+        assert!(functools_module_stub_source.contains("def update_wrapper("));
+        Ok(())
+    }
+
     #[test]
     fn resolve_package() -> anyhow::Result<()> {
         let TestCase {