ty_vendored 0.0.10

//! 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/ty_vendored/vendor/typeshed` change.

use std::fs::File;
use std::io::Write;
use std::path::Path;

use path_slash::PathExt;
use zip::CompressionMethod;
use zip::result::ZipResult;
use zip::write::{SimpleFileOptions, ZipWriter};

const TYPESHED_SOURCE_DIR: &str = "vendor/typeshed";
const TY_EXTENSIONS_STUBS: &[(&str, &str)] = &[
    (
        "ty_extensions/__init__.pyi",
        "stdlib/ty_extensions/__init__.pyi",
    ),
    (
        "ty_extensions/_internal.pyi",
        "stdlib/ty_extensions/_internal.pyi",
    ),
    (
        "ty_extensions/pydantic.pyi",
        "stdlib/ty_extensions/pydantic.pyi",
    ),
];
const TYPESHED_ZIP_LOCATION: &str = "/zipped_typeshed.zip";

/// Recursively zip the contents of the entire typeshed directory and patch typeshed
/// on the fly to include the `ty_extensions` module.
///
/// This routine is adapted from a recipe at
/// <https://github.com/zip-rs/zip-old/blob/5d0f198124946b7be4e5969719a7f29f363118cd/examples/write_dir.rs>
fn write_zipped_typeshed_to(writer: File) -> ZipResult<File> {
    let mut zip = ZipWriter::new(writer);

    // Use deflated compression for WASM builds because compiling `zstd-sys` requires clang
    // [source](https://github.com/gyscos/zstd-rs/wiki/Compile-for-WASM) which complicates the build
    // by a lot. Deflated compression is slower but it shouldn't matter much for the WASM use case
    // (WASM itself is already slower than a native build for a specific platform).
    #[cfg(feature = "zstd")]
    let method = CompressionMethod::Zstd;
    #[cfg(all(not(feature = "zstd"), feature = "deflate"))]
    let method = CompressionMethod::Deflated;
    #[cfg(not(any(feature = "zstd", feature = "deflate")))]
    let method = CompressionMethod::Stored;

    let options = SimpleFileOptions::default()
        .compression_method(method)
        .unix_permissions(0o644);

    for entry in walkdir::WalkDir::new(TYPESHED_SOURCE_DIR) {
        let dir_entry = entry.unwrap();
        let absolute_path = dir_entry.path();
        let normalized_relative_path = absolute_path
            .strip_prefix(Path::new(TYPESHED_SOURCE_DIR))
            .unwrap()
            .to_slash()
            .expect("Unexpected non-utf8 typeshed path!");

        // Write file or directory explicitly
        // Some unzip tools unzip files with directory paths correctly, some do not!
        if absolute_path.is_file() {
            println!("adding file {absolute_path:?} as {normalized_relative_path:?} ...");
            zip.start_file(&*normalized_relative_path, options)?;
            let mut f = File::open(absolute_path)?;
            std::io::copy(&mut f, &mut zip).unwrap();

            // Patch the VERSIONS file to make `ty_extensions` available
            if normalized_relative_path == "stdlib/VERSIONS" {
                writeln!(&mut zip, "ty_extensions: 3.0-")?;
            }
        } else if !normalized_relative_path.is_empty() {
            // Only if not root! Avoids path spec / warning
            // and mapname conversion failed error on unzip
            println!("adding dir {absolute_path:?} as {normalized_relative_path:?} ...");
            zip.add_directory(normalized_relative_path, options)?;
        }
    }

    // Patch typeshed and add the stubs for the `ty_extensions` package.
    zip.add_directory("stdlib/ty_extensions/", options)?;
    for (source, destination) in TY_EXTENSIONS_STUBS {
        println!("adding file {source} as {destination} ...");
        zip.start_file(destination, options)?;
        let mut f = File::open(source)?;
        std::io::copy(&mut f, &mut zip).unwrap();
    }

    zip.finish()
}

fn main() {
    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 = File::create(zipped_typeshed_location).unwrap();
    write_zipped_typeshed_to(zipped_typeshed_file).unwrap();
}