From e824fe6d2b39365fdcacc416a721afffffa81035 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 6 Oct 2023 21:38:38 -0400 Subject: [PATCH] Copy over `install-wheel-rs` crate (#33) This PR copies over the `install-wheel-rs` crate at commit `10730ea1a84c58af6b35fb74c89ed0578ab042b6` with no modifications. It won't pass CI, but modifications will intentionally be confined to later PRs. --- crates/install-wheel-rs/Cargo.toml | 53 + crates/install-wheel-rs/Readme.md | 14 + crates/install-wheel-rs/install_wheel_rs.pyi | 3 + crates/install-wheel-rs/pyproject.toml | 9 + crates/install-wheel-rs/run_tests.sh | 12 + .../install-wheel-rs/src/install_location.rs | 188 +++ crates/install-wheel-rs/src/lib.rs | 123 ++ crates/install-wheel-rs/src/main.rs | 81 + crates/install-wheel-rs/src/pip_compileall.py | 31 + .../install-wheel-rs/src/python_bindings.rs | 93 ++ crates/install-wheel-rs/src/wheel.rs | 1455 +++++++++++++++++ crates/install-wheel-rs/src/wheel_tags.rs | 937 +++++++++++ .../upsidedown-0.4-py2.py3-none-any.whl | Bin 0 -> 5495 bytes .../test/test_install_wheel_rs.py | 35 + .../install-wheel-rs/windows-launcher/t32.exe | Bin 0 -> 171854 bytes .../windows-launcher/t64-arm.exe | Bin 0 -> 292806 bytes .../install-wheel-rs/windows-launcher/t64.exe | Bin 0 -> 180211 bytes 17 files changed, 3034 insertions(+) create mode 100644 crates/install-wheel-rs/Cargo.toml create mode 100644 crates/install-wheel-rs/Readme.md create mode 100644 crates/install-wheel-rs/install_wheel_rs.pyi create mode 100644 crates/install-wheel-rs/pyproject.toml create mode 100644 crates/install-wheel-rs/run_tests.sh create mode 100644 crates/install-wheel-rs/src/install_location.rs create mode 100644 crates/install-wheel-rs/src/lib.rs create mode 100644 crates/install-wheel-rs/src/main.rs create mode 100644 crates/install-wheel-rs/src/pip_compileall.py create mode 100644 crates/install-wheel-rs/src/python_bindings.rs create mode 100644 crates/install-wheel-rs/src/wheel.rs create mode 100644 crates/install-wheel-rs/src/wheel_tags.rs create mode 100644 crates/install-wheel-rs/test-data/upsidedown-0.4-py2.py3-none-any.whl create mode 100644 crates/install-wheel-rs/test/test_install_wheel_rs.py create mode 100644 crates/install-wheel-rs/windows-launcher/t32.exe create mode 100644 crates/install-wheel-rs/windows-launcher/t64-arm.exe create mode 100644 crates/install-wheel-rs/windows-launcher/t64.exe diff --git a/crates/install-wheel-rs/Cargo.toml b/crates/install-wheel-rs/Cargo.toml new file mode 100644 index 000000000..69c00cdf7 --- /dev/null +++ b/crates/install-wheel-rs/Cargo.toml @@ -0,0 +1,53 @@ +[package] +name = "install-wheel-rs" +version = "0.0.1" +edition = "2021" +description = "Takes a wheel and installs it, either in a venv or for monotrail" +license = "MIT OR Apache-2.0" +repository = "https://github.com/konstin/poc-monotrail" +keywords = ["wheel", "python"] + +[lib] +name = "install_wheel_rs" +# https://github.com/PyO3/maturin/issues/1080 :(( +#crate-type = ["cdylib", "rlib"] + +[dependencies] +clap = { version = "4.4.6", optional = true, features = ["derive", "env"] } +configparser = "3.0.2" +csv = "1.2.2" +data-encoding = "2.4.0" +fs-err = { workspace = true } +fs2 = { workspace = true } +glibc_version = "0.1.2" +goblin = "0.7.1" +mailparse = "0.14.0" +once_cell = "1.18.0" +platform-info = "2.0.2" +plist = "1.5.0" +pyo3 = { workspace = true, features = ["extension-module", "abi3-py37"], optional = true } +rayon = { version = "1.8.0", optional = true } +regex = { workspace = true } +rfc2047-decoder = "1.0.1" +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +sha2 = { workspace = true } +target-lexicon = "0.12.11" +tempfile = { workspace = true } +thiserror = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true, optional = true } +walkdir = { workspace = true } +zip = { version = "0.6.6", default-features = false, features = ["deflate"] } # no default features for zstd + +[features] +default = ["cli", "parallel"] +python_bindings = ["pyo3", "tracing-subscriber"] +cli = ["clap"] +parallel = ["rayon"] + +[dev-dependencies] +indoc = { workspace = true } + +[package.metadata.dist] +dist = false diff --git a/crates/install-wheel-rs/Readme.md b/crates/install-wheel-rs/Readme.md new file mode 100644 index 000000000..01cb6624a --- /dev/null +++ b/crates/install-wheel-rs/Readme.md @@ -0,0 +1,14 @@ +Reimplementation of wheel installing in rust. Supports both classical venvs and monotrail. + +There are simple python bindings: + +```python +from install_wheel_rs import LockedVenv + +locked_venv = LockedVenv("path/to/.venv") +locked_venv.install_wheel("path/to/some_tagged_wheel.whl") +``` + +and there's only one function: `install_wheels_venv(wheels: List[str], venv: str)`, where `wheels` is a list of paths to wheel files and `venv` is the location of the venv to install the packages in. + +See monotrail for benchmarks. diff --git a/crates/install-wheel-rs/install_wheel_rs.pyi b/crates/install-wheel-rs/install_wheel_rs.pyi new file mode 100644 index 000000000..c2214beef --- /dev/null +++ b/crates/install-wheel-rs/install_wheel_rs.pyi @@ -0,0 +1,3 @@ +class LockedVenv: + def __init__(self, venv: str): ... + def install_wheel(self, wheel: str): ... diff --git a/crates/install-wheel-rs/pyproject.toml b/crates/install-wheel-rs/pyproject.toml new file mode 100644 index 000000000..b81f0f237 --- /dev/null +++ b/crates/install-wheel-rs/pyproject.toml @@ -0,0 +1,9 @@ +[project] +name = "install-wheel-rs" + +[tool.maturin] +features = ["python_bindings"] + +[build-system] +requires = ["maturin>=1.2,<2.0"] +build-backend = "maturin" diff --git a/crates/install-wheel-rs/run_tests.sh b/crates/install-wheel-rs/run_tests.sh new file mode 100644 index 000000000..7b57c4fe8 --- /dev/null +++ b/crates/install-wheel-rs/run_tests.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +set -e + +# cd to project root +cd "$(git rev-parse --show-toplevel)" +rm -f target-maturin/wheels/install_wheel_rs-*.whl +CARGO_TARGET_DIR=target-maturin maturin build --release --strip --no-sdist -m install-wheel-rs/Cargo.toml +.venv/bin/pip uninstall -y install-wheel-rs +.venv/bin/pip install target-maturin/wheels/install_wheel_rs-*.whl +.venv/bin/pytest install-wheel-rs/test +cargo test \ No newline at end of file diff --git a/crates/install-wheel-rs/src/install_location.rs b/crates/install-wheel-rs/src/install_location.rs new file mode 100644 index 000000000..29fade9bd --- /dev/null +++ b/crates/install-wheel-rs/src/install_location.rs @@ -0,0 +1,188 @@ +//! Multiplexing between venv install and monotrail install + +use fs2::FileExt; +use fs_err as fs; +use fs_err::File; +use std::io; +use std::ops::Deref; +use std::path::{Path, PathBuf}; +use tracing::{error, warn}; + +const INSTALL_LOCKFILE: &str = "install-wheel-rs.lock"; + +/// I'm not sure that's the right way to normalize here, but it's a single place to change +/// everything. +/// +/// For displaying to the user, `-` is better, and it's also what poetry lockfile 2.0 does +/// +/// Keep in sync with `find_distributions` +pub fn normalize_name(dep_name: &str) -> String { + dep_name.to_lowercase().replace(['.', '_'], "-") +} + +/// A directory for which we acquired a install-wheel-rs.lock lockfile +pub struct LockedDir { + /// The directory to lock + path: PathBuf, + /// handle on the install-wheel-rs.lock that drops the lock + lockfile: File, +} + +impl LockedDir { + /// Tries to lock the directory, returns Ok(None) if it is already locked + pub fn try_acquire(path: &Path) -> io::Result> { + let lockfile = File::create(path.join(INSTALL_LOCKFILE))?; + if lockfile.file().try_lock_exclusive().is_ok() { + Ok(Some(Self { + path: path.to_path_buf(), + lockfile, + })) + } else { + Ok(None) + } + } + + /// Locks the directory, if necessary blocking until the lock becomes free + pub fn acquire(path: &Path) -> io::Result { + let lockfile = File::create(path.join(INSTALL_LOCKFILE))?; + lockfile.file().lock_exclusive()?; + Ok(Self { + path: path.to_path_buf(), + lockfile, + }) + } +} + +impl Drop for LockedDir { + fn drop(&mut self) { + if let Err(err) = self.lockfile.file().unlock() { + error!( + "Failed to unlock {}: {}", + self.lockfile.path().display(), + err + ); + } + } +} + +impl Deref for LockedDir { + type Target = Path; + + fn deref(&self) -> &Self::Target { + &self.path + } +} + +/// Multiplexing between venv install and monotrail install +/// +/// For monotrail, we have a structure that is {monotrail}/{normalized(name)}/{version}/tag +/// +/// We use a lockfile to prevent multiple instance writing stuff on the same time +/// As of pip 22.0, e.g. `pip install numpy; pip install numpy; pip install numpy` will +/// nondeterministically fail +/// +/// I was also thinking about making a shared lock on the import side, but monotrail install +/// is supposedly atomic (by directory renaming), while for venv installation there can't be +/// atomicity (we need to add lots of different file without a top level directory / key-turn +/// file we could rename) and the locking would also need to happen in the import mechanism +/// itself to ensure +pub enum InstallLocation> { + Venv { + /// absolute path + venv_base: T, + python_version: (u8, u8), + }, + Monotrail { + monotrail_root: T, + python: PathBuf, + python_version: (u8, u8), + }, +} + +impl> InstallLocation { + /// Returns the location of the python interpreter + pub fn get_python(&self) -> PathBuf { + match self { + InstallLocation::Venv { venv_base, .. } => { + if cfg!(windows) { + venv_base.join("Scripts").join("python.exe") + } else { + // canonicalize on python would resolve the symlink + venv_base.join("bin").join("python") + } + } + // TODO: For monotrail use the monotrail launcher + InstallLocation::Monotrail { python, .. } => python.clone(), + } + } + + pub fn get_python_version(&self) -> (u8, u8) { + match self { + InstallLocation::Venv { python_version, .. } => *python_version, + InstallLocation::Monotrail { python_version, .. } => *python_version, + } + } + + /// TODO: This function is unused? + pub fn is_installed(&self, normalized_name: &str, version: &str) -> bool { + match self { + InstallLocation::Venv { + venv_base, + python_version, + } => { + let site_packages = if cfg!(target_os = "windows") { + venv_base.join("Lib").join("site-packages") + } else { + venv_base + .join("lib") + .join(format!("python{}.{}", python_version.0, python_version.1)) + .join("site-packages") + }; + site_packages + .join(format!("{}-{}.dist-info", normalized_name, version)) + .is_dir() + } + InstallLocation::Monotrail { monotrail_root, .. } => monotrail_root + .join(format!("{}-{}", normalized_name, version)) + .is_dir(), + } + } +} + +impl InstallLocation { + pub fn acquire_lock(&self) -> io::Result> { + let root = match self { + Self::Venv { venv_base, .. } => venv_base, + Self::Monotrail { monotrail_root, .. } => monotrail_root, + }; + + // If necessary, create monotrail dir + fs::create_dir_all(root)?; + + let locked_dir = if let Some(locked_dir) = LockedDir::try_acquire(root)? { + locked_dir + } else { + warn!( + "Could not acquire exclusive lock for installing, is another installation process \ + running? Sleeping until lock becomes free" + ); + LockedDir::acquire(root)? + }; + + Ok(match self { + Self::Venv { python_version, .. } => InstallLocation::Venv { + venv_base: locked_dir, + python_version: *python_version, + }, + Self::Monotrail { + python_version, + python, + .. + } => InstallLocation::Monotrail { + monotrail_root: locked_dir, + python: python.clone(), + python_version: *python_version, + }, + }) + } +} diff --git a/crates/install-wheel-rs/src/lib.rs b/crates/install-wheel-rs/src/lib.rs new file mode 100644 index 000000000..c73020d12 --- /dev/null +++ b/crates/install-wheel-rs/src/lib.rs @@ -0,0 +1,123 @@ +//! Takes a wheel and installs it, either in a venv or for monotrail +//! +//! ```no_run +//! use std::path::Path; +//! use install_wheel_rs::install_wheel_in_venv; +//! +//! install_wheel_in_venv( +//! "Django-4.2.6-py3-none-any.whl", +//! ".venv", +//! ".venv/bin/python", +//! (3, 8), +//! ).unwrap(); +//! ``` + +use platform_info::PlatformInfoError; +use std::fs::File; +use std::io; +use std::path::Path; +use std::str::FromStr; +use thiserror::Error; +use zip::result::ZipError; + +pub use install_location::{normalize_name, InstallLocation, LockedDir}; +pub use wheel::{ + get_script_launcher, install_wheel, parse_key_value_file, read_record_file, relative_to, + Script, SHEBANG_PYTHON, +}; +pub use wheel_tags::{Arch, CompatibleTags, Os, WheelFilename}; + +mod install_location; +#[cfg(feature = "python_bindings")] +mod python_bindings; +mod wheel; +mod wheel_tags; + +#[derive(Error, Debug)] +pub enum Error { + #[error(transparent)] + IO(#[from] io::Error), + /// This shouldn't actually be possible to occur + #[error("Failed to serialize direct_url.json ಠ_ಠ")] + DirectUrlSerdeJson(#[source] serde_json::Error), + /// Tags/metadata didn't match platform + #[error("The wheel is incompatible with the current platform {os} {arch}")] + IncompatibleWheel { os: Os, arch: Arch }, + /// The wheel is broken + #[error("The wheel is invalid: {0}")] + InvalidWheel(String), + /// pyproject.toml or poetry.lock are broken + #[error("The poetry dependency specification (pyproject.toml or poetry.lock) is broken (try `poetry update`?): {0}")] + InvalidPoetry(String), + /// Doesn't follow file name schema + #[error("The wheel filename \"{0}\" is invalid: {1}")] + InvalidWheelFileName(String, String), + #[error("Failed to read the wheel file {0}")] + Zip(String, #[source] ZipError), + #[error("Failed to run python subcommand")] + PythonSubcommand(#[source] io::Error), + #[error("Failed to move data files")] + WalkDir(#[from] walkdir::Error), + #[error("RECORD file doesn't match wheel contents: {0}")] + RecordFile(String), + #[error("RECORD file is invalid")] + RecordCsv(#[from] csv::Error), + #[error("Broken virtualenv: {0}")] + BrokenVenv(String), + #[error("Failed to detect the operating system version: {0}")] + OsVersionDetection(String), + #[error("Failed to detect the current platform")] + PlatformInfo(#[source] PlatformInfoError), + #[error("Invalid version specification, only none or == is supported")] + Pep440, +} + +impl Error { + pub(crate) fn from_zip_error(file: String, value: ZipError) -> Self { + match value { + ZipError::Io(io_error) => Self::IO(io_error), + _ => Self::Zip(file, value), + } + } +} + +/// High level API: Install a wheel in a virtualenv +/// +/// The python interpreter is used for compiling to byte code, the python version for computing +/// the site packages path on unix. +/// +/// Returns the tag of the wheel +pub fn install_wheel_in_venv( + wheel: impl AsRef, + venv: impl AsRef, + interpreter: impl AsRef, + major_minor: (u8, u8), +) -> Result { + let venv_base = venv.as_ref().canonicalize()?; + let location = InstallLocation::Venv { + venv_base, + python_version: major_minor, + }; + let locked_dir = location.acquire_lock()?; + + let filename = wheel + .as_ref() + .file_name() + .ok_or_else(|| Error::InvalidWheel("Expected a file".to_string()))? + .to_string_lossy(); + let filename = WheelFilename::from_str(&filename)?; + let compatible_tags = CompatibleTags::current(location.get_python_version())?; + filename.compatibility(&compatible_tags)?; + + install_wheel( + &locked_dir, + File::open(wheel)?, + filename, + false, + true, + &[], + // Only relevant for monotrail style installation + "", + interpreter, + ) +} diff --git a/crates/install-wheel-rs/src/main.rs b/crates/install-wheel-rs/src/main.rs new file mode 100644 index 000000000..e45183e75 --- /dev/null +++ b/crates/install-wheel-rs/src/main.rs @@ -0,0 +1,81 @@ +use clap::Parser; +use fs_err::File; +use install_wheel_rs::{install_wheel, CompatibleTags, Error, InstallLocation, WheelFilename}; +#[cfg(feature = "rayon")] +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use std::path::PathBuf; +use std::str::FromStr; + +/// Low level install CLI, mainly used for testing +#[derive(Parser)] +struct Args { + wheels: Vec, + /// The root of the venv to install in + #[clap(long, env = "VIRTUAL_ENV")] + venv: PathBuf, + /// The major version of the current python interpreter + #[clap(long)] + major: u8, + /// The minor version of the current python interpreter + #[clap(long)] + minor: u8, + /// Compile .py files to .pyc (errors are ignored) + #[clap(long)] + compile: bool, + /// Don't check the hashes in RECORD + #[clap(long)] + skip_hashes: bool, +} + +fn main() -> Result<(), Error> { + let args = Args::parse(); + let venv_base = args.venv.canonicalize()?; + let location = InstallLocation::Venv { + venv_base, + python_version: (args.major, args.minor), + }; + let locked_dir = location.acquire_lock()?; + + let wheels: Vec<(PathBuf, WheelFilename)> = args + .wheels + .into_iter() + .map(|wheel| { + let filename = wheel + .file_name() + .ok_or_else(|| Error::InvalidWheel("Expected a file".to_string()))? + .to_string_lossy(); + let filename = WheelFilename::from_str(&filename)?; + let compatible_tags = CompatibleTags::current(location.get_python_version())?; + filename.compatibility(&compatible_tags)?; + Ok((wheel, filename)) + }) + .collect::>()?; + + let wheels = { + #[cfg(feature = "rayon")] + { + wheels.into_par_iter() + } + #[cfg(not(feature = "rayon"))] + { + wheels.into_iter() + } + }; + wheels + .map(|(wheel, filename)| { + install_wheel( + &locked_dir, + File::open(wheel)?, + filename, + args.compile, + !args.skip_hashes, + &[], + // Only relevant for monotrail style installation + "", + location.get_python(), + )?; + Ok(()) + }) + .collect::>()?; + Ok(()) +} diff --git a/crates/install-wheel-rs/src/pip_compileall.py b/crates/install-wheel-rs/src/pip_compileall.py new file mode 100644 index 000000000..bd6b79ea2 --- /dev/null +++ b/crates/install-wheel-rs/src/pip_compileall.py @@ -0,0 +1,31 @@ +""" +Based on +https://github.com/pypa/pip/blob/3820b0e52c7fed2b2c43ba731b718f316e6816d1/src/pip/_internal/operations/install/wheel.py#L612-L623 + +pip silently just swallows all pyc compilation errors, but `python -m compileall` does +not have such a flag, so we adapt the pip code. This is relevant e.g. for +`debugpy-1.5.1-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64`, +which contains some vendored python 2 code which fails to compile +""" + +import compileall +import sys +import warnings + +with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + # in rust, we give one line per file to compile + # we also have to read it before printing to stdout or we risk pipes running full + paths = sys.stdin.readlines() + for path in paths: + # just to be sure + path = path.strip() + if not path: + continue + # Unlike pip, we set quiet=2, so we don't have to capture stdout + # I'd like to show those errors, but given that pip thinks that's totally fine + # we can't really change that + success = compileall.compile_file(path, force=True, quiet=2) + if success: + # return successfully compiled files so we can update RECORD accordingly + print(path) diff --git a/crates/install-wheel-rs/src/python_bindings.rs b/crates/install-wheel-rs/src/python_bindings.rs new file mode 100644 index 000000000..edaa20684 --- /dev/null +++ b/crates/install-wheel-rs/src/python_bindings.rs @@ -0,0 +1,93 @@ +#![allow(clippy::format_push_string)] // I will not replace clear and infallible with fallible, io looking code + +use crate::{install_wheel, CompatibleTags, Error, InstallLocation, LockedDir, WheelFilename}; +use pyo3::create_exception; +use pyo3::types::PyModule; +use pyo3::{pyclass, pymethods, pymodule, PyErr, PyResult, Python}; +use std::env; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::str::FromStr; + +create_exception!( + install_wheel_rs, + PyWheelInstallerError, + pyo3::exceptions::PyException +); + +impl From for PyErr { + fn from(err: Error) -> Self { + let mut accumulator = format!("Failed to install wheels: {}", err); + + let mut current_err: &dyn std::error::Error = &err; + while let Some(cause) = current_err.source() { + accumulator.push_str(&format!("\n Caused by: {}", cause)); + current_err = cause; + } + PyWheelInstallerError::new_err(accumulator) + } +} + +#[pyclass] +struct LockedVenv { + location: InstallLocation, +} + +#[pymethods] +impl LockedVenv { + #[new] + pub fn new(py: Python, venv: PathBuf) -> PyResult { + Ok(Self { + location: InstallLocation::Venv { + venv_base: LockedDir::acquire(&venv)?, + python_version: (py.version_info().major, py.version_info().minor), + }, + }) + } + + pub fn install_wheel(&self, py: Python, wheel: PathBuf) -> PyResult<()> { + // Would be nicer through https://docs.python.org/3/c-api/init.html#c.Py_GetProgramFullPath + let sys_executable: String = py.import("sys")?.getattr("executable")?.extract()?; + + // TODO: Pass those options on to the user + py.allow_threads(|| { + let filename = wheel + .file_name() + .ok_or_else(|| Error::InvalidWheel("Expected a file".to_string()))? + .to_string_lossy(); + let filename = WheelFilename::from_str(&filename)?; + let compatible_tags = CompatibleTags::current(self.location.get_python_version())?; + filename.compatibility(&compatible_tags)?; + + install_wheel( + &self.location, + File::open(wheel)?, + filename, + true, + true, + &[], + // unique_version can be anything since it's only used to monotrail + "", + Path::new(&sys_executable), + ) + })?; + Ok(()) + } +} + +#[pymodule] +pub fn install_wheel_rs(_py: Python, m: &PyModule) -> PyResult<()> { + // Good enough for now + if env::var_os("RUST_LOG").is_some() { + tracing_subscriber::fmt::init(); + } else { + let format = tracing_subscriber::fmt::format() + .with_level(false) + .with_target(false) + .without_time() + .compact(); + tracing_subscriber::fmt().event_format(format).init(); + } + m.add_class::()?; + Ok(()) +} diff --git a/crates/install-wheel-rs/src/wheel.rs b/crates/install-wheel-rs/src/wheel.rs new file mode 100644 index 000000000..f4274d6bf --- /dev/null +++ b/crates/install-wheel-rs/src/wheel.rs @@ -0,0 +1,1455 @@ +#![allow(clippy::needless_borrow)] + +use crate::install_location::{InstallLocation, LockedDir}; +use crate::wheel_tags::WheelFilename; +use crate::{normalize_name, Error}; +use configparser::ini::Ini; +use data_encoding::BASE64URL_NOPAD; +use fs_err as fs; +use fs_err::{DirEntry, File}; +use mailparse::MailHeaderMap; +use regex::Regex; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::collections::{HashMap, HashSet}; +use std::ffi::OsString; +use std::io::{BufRead, BufReader, BufWriter, Cursor, Read, Seek, Write}; +use std::path::{Path, PathBuf}; +use std::process::{Command, ExitStatus, Stdio}; +use std::{env, io, iter}; +use tempfile::{tempdir, TempDir}; +use tracing::{debug, error, span, warn, Level}; +use walkdir::WalkDir; +use zip::result::ZipError; +use zip::write::FileOptions; +use zip::{ZipArchive, ZipWriter}; + +/// `#!/usr/bin/env python` +pub const SHEBANG_PYTHON: &str = "#!/usr/bin/env python"; + +pub const LAUNCHER_T32: &[u8] = include_bytes!("../windows-launcher/t32.exe"); +pub const LAUNCHER_T64: &[u8] = include_bytes!("../windows-launcher/t64.exe"); +pub const LAUNCHER_T64_ARM: &[u8] = include_bytes!("../windows-launcher/t64-arm.exe"); + +/// Line in a RECORD file +/// +/// +/// ```csv +/// tqdm/cli.py,sha256=x_c8nmc4Huc-lKEsAXj78ZiyqSJ9hJ71j7vltY67icw,10509 +/// tqdm-4.62.3.dist-info/RECORD,, +/// ``` +#[derive(Deserialize, Serialize, PartialOrd, PartialEq, Ord, Eq)] +pub struct RecordEntry { + pub path: String, + pub hash: Option, + #[allow(dead_code)] + pub size: Option, +} + +/// Minimal direct_url.json schema +/// +/// +/// +#[derive(Serialize)] +struct DirectUrl { + archive_info: HashMap<(), ()>, + url: String, +} + +/// A script defining the name of the runnable entrypoint and the module and function that should be +/// run. +#[cfg(feature = "python_bindings")] +#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +#[pyo3::pyclass(dict)] +pub struct Script { + #[pyo3(get)] + pub script_name: String, + #[pyo3(get)] + pub module: String, + #[pyo3(get)] + pub function: String, +} + +/// A script defining the name of the runnable entrypoint and the module and function that should be +/// run. +#[cfg(not(feature = "python_bindings"))] +#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +pub struct Script { + pub script_name: String, + pub module: String, + pub function: String, +} + +impl Script { + /// Parses a script definition like `foo.bar:main` or `foomod:main_bar [bar,baz]` + /// + /// + /// + /// Extras are supposed to be ignored, which happens if you pass None for extras + pub fn from_value( + script_name: &str, + value: &str, + extras: Option<&[String]>, + ) -> Result, Error> { + let script_regex = Regex::new(r"^(?P[\w\d_\-.]+):(?P[\w\d_\-.]+)(?:\s+\[(?P(?:[^,]+,?\s*)+)\])?$").unwrap(); + + let captures = script_regex + .captures(value) + .ok_or_else(|| Error::InvalidWheel(format!("invalid console script: '{}'", value)))?; + if let Some(script_extras) = captures.name("extras") { + let script_extras = script_extras + .as_str() + .split(',') + .map(|extra| extra.trim().to_string()) + .collect::>(); + if let Some(extras) = extras { + if !script_extras.is_subset(&extras.iter().cloned().collect()) { + return Ok(None); + } + } + } + + Ok(Some(Script { + script_name: script_name.to_string(), + module: captures.name("module").unwrap().as_str().to_string(), + function: captures.name("function").unwrap().as_str().to_string(), + })) + } +} + +/// Wrapper script template function +/// +/// +pub fn get_script_launcher(module: &str, import_name: &str, shebang: &str) -> String { + format!( + r##"{shebang} +# -*- coding: utf-8 -*- +import re +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}()) +"##, + shebang = shebang, + module = module, + import_name = import_name + ) +} + +/// Part of entrypoints parsing +fn read_scripts_from_section( + scripts_section: &HashMap>, + section_name: &str, + extras: Option<&[String]>, +) -> Result, Error> { + let mut scripts = Vec::new(); + for (script_name, python_location) in scripts_section.iter() { + match python_location { + Some(value) => { + if let Some(script) = Script::from_value(script_name, value, extras)? { + scripts.push(script); + } + } + None => { + return Err(Error::InvalidWheel(format!( + "[{}] key {} must have a value", + section_name, script_name + ))); + } + } + } + Ok(scripts) +} + +/// Parses the entry_points.txt entry in the wheel for console scripts +/// +/// Returns (script_name, module, function) +/// +/// Extras are supposed to be ignored, which happens if you pass None for extras +fn parse_scripts( + archive: &mut ZipArchive, + dist_info_prefix: &str, + extras: Option<&[String]>, +) -> Result<(Vec