From 4d6e953a34e5acc933da8bf4620e32194d66dfab Mon Sep 17 00:00:00 2001 From: xenia Date: Wed, 16 Apr 2025 18:32:44 -0400 Subject: [PATCH] add ipython magic importer --- python/nix-import/51-nix-importer.py | 170 +++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 python/nix-import/51-nix-importer.py diff --git a/python/nix-import/51-nix-importer.py b/python/nix-import/51-nix-importer.py new file mode 100644 index 0000000..380eaea --- /dev/null +++ b/python/nix-import/51-nix-importer.py @@ -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 +# 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(" 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