add ipython magic importer
This commit is contained in:
parent
3b78341c40
commit
4d6e953a34
|
|
@ -0,0 +1,170 @@
|
|||
# Sets up interactive sessions in ipython to automatically import missing modules using nix Depends
|
||||
# on nix-index/nix-locate being installed and functional, see
|
||||
# <https://github.com/nix-community/nix-index-database> for a prebuilt database.
|
||||
#
|
||||
# How it works:
|
||||
# - During startup, the sysconfig is used to determine where to find site-packages and the probable
|
||||
# nixpkgs variable name for the python scope.
|
||||
# - To find a package, nix-locate is used to try to find a package which exports
|
||||
# {sitePackages}/{name}/__init__.py. If that is found, it attempts to download nixpkgs#{result}.
|
||||
# - If that fails, it instead attempts to build nixpkgs#{pythonScope}.{name}.
|
||||
# - When either of those strategies result in nix store paths, the store paths are recursively
|
||||
# processed (taking propagated-build-inputs into account) and any {sitePackages} found in any
|
||||
# dependency is added to sys.path.
|
||||
# - Finally, the import is tried again with importlib.machinery.PathFinder
|
||||
# - A directory is created under $XDG_RUNTIME_DIR/ipython per process which contains GC roots for
|
||||
# any packages loaded this way. When ipython exits, the per-process directory is deleted. This
|
||||
# ensures that a GC during a running ipython session doesn't nuke any temporarily installed
|
||||
# packages.
|
||||
#
|
||||
# Limitations:
|
||||
# - The only way to be able to look up a package that is not named the same as its top-level module
|
||||
# is using nix-locate. The index is only built for the default python scope, and not any other
|
||||
# versions of python. This limitation could be worked around by scanning for pythonImportsCheck in
|
||||
# nixpkgs but this is not performed currently. Therefore, eg `import Crypto` works on the default
|
||||
# python version for a given NixOS installation, but not any other version. This script also
|
||||
# provids install_with_nix(package), which can be called with eg install_with_nix("pycryptodome"),
|
||||
# after which such imports should work.
|
||||
# - Support for propagated-build-inputs is very rudimentary. For example, a python package which
|
||||
# depends on a binary being present in PATH which is declared in propagatedBuildInputs would not
|
||||
# currently work. However, pure-python dependencies work as expected.
|
||||
# - This will probably explode and should not be used outside of **interactive** sessions, there is
|
||||
# a check to make sure an ipython input is in the call stack when the MetaPathFinder is triggered.
|
||||
|
||||
|
||||
def _nix_importer_init():
|
||||
import atexit
|
||||
import importlib.abc
|
||||
import importlib.machinery
|
||||
import importlib.util
|
||||
import inspect
|
||||
import os
|
||||
import pathlib
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import sysconfig
|
||||
import tempfile
|
||||
from types import ModuleType
|
||||
from typing import Sequence
|
||||
|
||||
|
||||
runtime_dir = pathlib.Path(
|
||||
os.environ.get("XDG_RUNTIME_DIR",
|
||||
str(pathlib.Path("/run/user") / str(os.getuid()))))
|
||||
ipython_dir = runtime_dir / "ipython"
|
||||
if not ipython_dir.is_dir():
|
||||
ipython_dir.mkdir()
|
||||
session_dir = ipython_dir / f"session-{os.getpid()}"
|
||||
if session_dir.exists():
|
||||
shutil.rmtree(session_dir)
|
||||
session_dir.mkdir()
|
||||
atexit.register(lambda: shutil.rmtree(session_dir))
|
||||
|
||||
|
||||
class NixImporter(importlib.abc.MetaPathFinder):
|
||||
def __init__(self) -> None:
|
||||
self._result_counter = 0
|
||||
self._subfinder = importlib.machinery.PathFinder()
|
||||
|
||||
config_vars = sysconfig.get_config_vars()
|
||||
config_vars["base"] = ""
|
||||
path = sysconfig.get_path("purelib", scheme="posix_prefix", vars=config_vars)
|
||||
# remove leading / if it exists
|
||||
self._base_path = ("/" / pathlib.Path(path)).relative_to("/")
|
||||
self._base_packages = f"python{config_vars['py_version_nodot']}Packages"
|
||||
|
||||
|
||||
def find_spec(self, fullname: str,
|
||||
path: Sequence[str] | None,
|
||||
target: None | ModuleType = None) -> importlib.machinery.ModuleSpec | None:
|
||||
# Deny autoimport to the nix store, and only allow it if the traceback indicates ipython
|
||||
f = inspect.currentframe()
|
||||
while f is not None:
|
||||
if f.f_code.co_filename.startswith("/nix/store"):
|
||||
return None
|
||||
elif f.f_code.co_filename.startswith("<ipython-input"):
|
||||
break
|
||||
f = f.f_back
|
||||
else:
|
||||
return None
|
||||
|
||||
find_path = self._base_path
|
||||
for element in fullname.split("."):
|
||||
find_path = find_path / element
|
||||
find_path = find_path / "__init__.py"
|
||||
|
||||
print(f"Module {fullname} not found, searching NixOS...")
|
||||
command = ["nix-locate", "-w", "--minimal", str(find_path)]
|
||||
|
||||
selected_target: str | None = None
|
||||
|
||||
try:
|
||||
output = subprocess.check_output(command).strip().decode()
|
||||
if len(output) > 0:
|
||||
if len(matches := output.split("\n")) > 1:
|
||||
print("Multiple matches found!", matches)
|
||||
return None
|
||||
else:
|
||||
selected_target = output
|
||||
except Exception:
|
||||
print("Failed to search")
|
||||
|
||||
if selected_target is None and "." not in fullname:
|
||||
selected_target = f"{self._base_packages}.{fullname}"
|
||||
|
||||
if selected_target is None:
|
||||
print("Could not find installation target!")
|
||||
return None
|
||||
|
||||
try:
|
||||
print("Installing", selected_target)
|
||||
self._install_with_nix(selected_target)
|
||||
return self._subfinder.find_spec(fullname, path, target)
|
||||
except Exception:
|
||||
print("Failed to install")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def invalidate_caches(self) -> None:
|
||||
self._subfinder.invalidate_caches()
|
||||
|
||||
def _process_propagated_build_inputs(self, path: pathlib.Path) -> None:
|
||||
python_path = path / self._base_path
|
||||
if python_path.exists() and str(python_path) not in sys.path:
|
||||
print("Appending path", str(python_path))
|
||||
sys.path.append(str(python_path))
|
||||
|
||||
nix_support = path / "nix-support"
|
||||
if not nix_support.is_dir():
|
||||
return
|
||||
propagated_build_inputs = nix_support / "propagated-build-inputs"
|
||||
if not propagated_build_inputs.is_file():
|
||||
return
|
||||
with propagated_build_inputs.open("r") as f:
|
||||
paths = f.read().strip().split(" ")
|
||||
for child_path in paths:
|
||||
self._process_propagated_build_inputs(pathlib.Path(child_path))
|
||||
|
||||
def _install_with_nix(self, path: str) -> None:
|
||||
if not path.startswith("/nix/store"):
|
||||
path = "nixpkgs#" + path
|
||||
|
||||
out_paths = subprocess.check_output(
|
||||
["nix", "build", "--out-link", f"link-{self._result_counter}", "--print-out-paths", path],
|
||||
cwd=session_dir).strip().decode().split("\n")
|
||||
self._result_counter += 1
|
||||
|
||||
for result in out_paths:
|
||||
self._process_propagated_build_inputs(pathlib.Path(result))
|
||||
|
||||
|
||||
importer = NixImporter()
|
||||
sys.meta_path.append(importer)
|
||||
|
||||
return importer
|
||||
|
||||
|
||||
install_with_nix = _nix_importer_init()._install_with_nix
|
||||
del _nix_importer_init
|
||||
Loading…
Reference in New Issue