add ipython magic importer

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
+# <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