From ac16e944914795cbc7df4fb9f3d2fae63e1da7d0 Mon Sep 17 00:00:00 2001 From: messense Date: Fri, 3 Dec 2021 22:45:30 +0800 Subject: [PATCH] Add a Python import hook --- Changelog.md | 1 + maturin/import_hook.py | 149 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 maturin/import_hook.py diff --git a/Changelog.md b/Changelog.md index 389dfbeac..ea5b19160 100644 --- a/Changelog.md +++ b/Changelog.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Don't package non-path-dep crates in sdist for workspaces in [#720](https://github.com/PyO3/maturin/pull/720) * Build release packages with `password-storage` feature in [#725](https://github.com/PyO3/maturin/pull/725) * Add support for x86_64 DargonFly BSD in [#727](https://github.com/PyO3/maturin/pull/727) +* Add a Python import hook in [#729](https://github.com/PyO3/maturin/pull/729) ## [0.12.3] - 2021-11-29 diff --git a/maturin/import_hook.py b/maturin/import_hook.py new file mode 100644 index 000000000..a989c511a --- /dev/null +++ b/maturin/import_hook.py @@ -0,0 +1,149 @@ +import contextlib +import importlib +import importlib.util +from importlib import abc +from importlib.machinery import ModuleSpec +import os +import pathlib +import shutil +import sys +import subprocess +from typing import Optional + +import toml + + +class Importer(abc.MetaPathFinder): + """A meta-path importer for the maturin based packages""" + + def __init__(self, bindings: Optional[str] = None, release: bool = False): + self.bindings = bindings + self.release = release + + def find_spec(self, fullname, path, target=None): + if fullname in sys.modules: + return + mod_parts = fullname.split(".") + module_name = mod_parts[-1] + + cwd = pathlib.Path(os.getcwd()) + # Full Cargo project in cwd + cargo_toml = cwd / "Cargo.toml" + if _is_cargo_project(cargo_toml, module_name): + return self._build_and_load(fullname, cargo_toml) + + # Full Cargo project in subdirectory of cwd + cargo_toml = cwd / module_name / "Cargo.toml" + if _is_cargo_project(cargo_toml, module_name): + return self._build_and_load(fullname, cargo_toml) + # module name with '-' instead of '_' + cargo_toml = cwd / module_name.replace("_", "-") / "Cargo.toml" + if _is_cargo_project(cargo_toml, module_name): + return self._build_and_load(fullname, cargo_toml) + + # Single .rs file + rust_file = cwd / (module_name + ".rs") + if rust_file.exists(): + project_dir = generate_project(rust_file, bindings=self.bindings or "pyo3") + cargo_toml = project_dir / "Cargo.toml" + return self._build_and_load(fullname, cargo_toml) + + def _build_and_load(self, fullname: str, cargo_toml: pathlib.Path) -> ModuleSpec: + build_module(cargo_toml, bindings=self.bindings) + loader = Loader(fullname) + return importlib.util.spec_from_loader(fullname, loader) + + +class Loader(abc.Loader): + def __init__(self, fullname): + self.fullname = fullname + + def load_module(self, fullname): + return importlib.import_module(self.fullname) + + +def _is_cargo_project(cargo_toml: pathlib.Path, module_name: str) -> bool: + with contextlib.suppress(FileNotFoundError): + with open(cargo_toml) as f: + cargo = toml.load(f) + package_name = cargo.get("package", {}).get("name") + if ( + package_name == module_name + or package_name.replace("-", "_") == module_name + ): + return True + return False + + +def generate_project(rust_file: pathlib.Path, bindings: str = "pyo3") -> pathlib.Path: + build_dir = pathlib.Path(os.getcwd()) / "build" + project_dir = build_dir / rust_file.stem + if project_dir.exists(): + shutil.rmtree(project_dir) + + command = ["maturin", "new", "-b", bindings, project_dir] + result = subprocess.run(command, stdout=subprocess.PIPE) + if result.returncode != 0: + sys.stderr.write( + f"Error: command {command} returned non-zero exit status {result.returncode}\n" + ) + raise ImportError("Failed to generate cargo project") + + with open(rust_file) as f: + lib_rs_content = f.read() + lib_rs = project_dir / "src" / "lib.rs" + with open(lib_rs, "w") as f: + f.write(lib_rs_content) + return project_dir + + +def build_module( + manifest_path: pathlib.Path, bindings: Optional[str] = None, release: bool = False +): + command = ["maturin", "develop", "-m", manifest_path] + if bindings: + command.append("-b") + command.append(bindings) + if release: + command.append("--release") + result = subprocess.run(command, stdout=subprocess.PIPE) + sys.stdout.buffer.write(result.stdout) + sys.stdout.flush() + if result.returncode != 0: + sys.stderr.write( + f"Error: command {command} returned non-zero exit status {result.returncode}\n" + ) + raise ImportError("Failed to build module with maturin") + + +def _have_importer() -> bool: + for importer in sys.meta_path: + if isinstance(importer, Importer): + return True + return False + + +def install(bindings: Optional[str] = None, release: bool = False): + """ + Install the import hook. + + :param bindings: Which kind of bindings to use. + Possible values are pyo3, rust-cpython and cffi + + :param release: Build in release mode, otherwise debug mode by default + """ + if _have_importer(): + return + importer = Importer(bindings=bindings, release=release) + sys.meta_path.append(importer) + return importer + + +def uninstall(importer: Importer): + """ + Uninstall the import hook. + """ + try: + sys.meta_path.remove(importer) + except ValueError: + pass