Skip to content

Commit

Permalink
Add a Python import hook
Browse files Browse the repository at this point in the history
  • Loading branch information
messense committed Dec 6, 2021
1 parent 7819bce commit a85e196
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 0 deletions.
1 change: 1 addition & 0 deletions Changelog.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
126 changes: 126 additions & 0 deletions maturin/import_hook.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit a85e196

Please sign in to comment.