171 lines
7.2 KiB
Python
171 lines
7.2 KiB
Python
# 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
|