From a85e196623a119d6bab23bcd6952f1ba6d450bdf Mon Sep 17 00:00:00 2001
From: messense <messense@icloud.com>
Date: Fri, 3 Dec 2021 22:45:30 +0800
Subject: [PATCH] Add a Python import hook

---
 Changelog.md           |   1 +
 maturin/import_hook.py | 126 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 127 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..eeb8ee126
--- /dev/null
+++ b/maturin/import_hook.py
@@ -0,0 +1,126 @@
+import importlib
+import importlib.util
+from importlib import abc
+import os
+import pathlib
+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]
+
+        # Full Cargo project
+        cargo_toml = pathlib.Path(os.getcwd()) / "Cargo.toml"
+        if os.path.exists(cargo_toml):
+            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
+                ):
+                    build_module(cargo_toml, bindings=self.bindings)
+                    loader = Loader(fullname)
+                    return importlib.util.spec_from_loader(fullname, loader)
+
+        # Single .rs file
+        rust_file = pathlib.Path(os.getcwd()) / (module_name + ".rs")
+        if os.path.exists(rust_file):
+            project_dir = generate_project(rust_file, bindings=self.bindings or "pyo3")
+            cargo_toml = project_dir / "Cargo.toml"
+            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 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
+    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