[red-knot] Include vendored typeshed stubs as a zipfile in the Ruff b…

…inary (#11779)

Co-authored-by: Micha Reiser <[email protected]>
Co-authored-by: Carl Meyer <[email protected]>

.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.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 {